import unittest

from utils import adquire_thread_orm
from plugins.disk_image.model import DiskProfile, Disk
from main.model import OrganizativeUnit, Computer
from pluginmanager.pluginmanager import PluginManager


class TestDP(unittest.TestCase):

    def test_models(self):
        orm = adquire_thread_orm()

        profile = DiskProfile()
        profile.ous.append(orm.query(OrganizativeUnit).filter_by(name=u'Facultad').one())
        profile.name = u"testfdp"
        profile.description = u"First disk profile description :D"
        disk = Disk()
        disk.add_partition(size=50, filesystem=u"NTFS", primary=True, bootable=True)
        disk.add_partition(size=10, filesystem=u"EXT3", primary=False)
        disk.add_partition(size=20, filesystem=u"SWAP", primary=False, bootable=True)
        disk.add_partition(size=10, primary=True)
        profile.add_disk(disk)
        orm.add(disk)
        orm.add(profile)
        orm.commit()

        profile = orm.query(DiskProfile).filter_by(name=u'testfdp').one()
        d = profile.get_disks()[0]
        partitions = d.get_partitions_order_by_num()
        p0 = partitions[0]
        # 1 is extended
        p1 = partitions[2]
        p2 = partitions[3]

        self.assertEqual(len(partitions), 4)
        self.assertEqual(p0.filesystem, u"NTFS")
        self.assertEqual(p2.filesystem, u"SWAP")
        self.assertEqual(p2.bootable, True)
        self.assertEqual(p1.bootable, False)
        self.assertEqual(p0.type, u"PRIMARY")
        self.assertEqual(p1.type, u"LOGIC")

    def test_getdp(self):

        orm = adquire_thread_orm()

        pm = PluginManager()
        plugin = pm.get_plugin_by_name('disk_image')

        profile = DiskProfile()
        profile.computers.append(orm.query(Computer).filter_by(name=u'ord1').one())
        profile.name = u"t1fdp"
        orm.add(profile)

        profile = DiskProfile()
        profile.ous.append(orm.query(OrganizativeUnit).filter_by(name=u'Aula A3.60').one())
        profile.name = u"t2fdp"
        orm.add(profile)
        orm.commit()

        c = orm.query(Computer).filter_by(name='fresa').one()
        dp = plugin.get_diskprofile_for_computer(c)
        self.assertEqual(dp.name, "t2fdp")

        c = orm.query(Computer).filter_by(name='ord1').one()
        dp = plugin.get_diskprofile_for_computer(c)
        self.assertEqual(dp.name, "t1fdp")


def suite():
   suite = unittest.TestSuite()
   suite.addTest(unittest.makeSuite(TestDP))
   return suite
