import unittest

from data_models.crate import Crate
from data_models.warehouse_inventory import WarehouseInventory
from data_models.workpiece import Workpiece, WorkpieceColor


class TestWarehouseInventory(unittest.TestCase):

    def test_constructor(self):
        """ Test constructor """
        wi = WarehouseInventory()
        self.assertEqual(len(wi.inventory), 9)
        for location, item in wi.inventory.items():
            self.assertEqual(item.workpiece, None)

    def test_insert(self):
        """Test insertion"""
        wi = WarehouseInventory()
        wi.take(wi.get_loc_empty_crate()) # remove a crate
        crate = Crate(Workpiece(0, WorkpieceColor.RED))
        crate_location = wi.insert(crate)
        self.assertEqual(crate, wi.get(crate_location))

    def test_get_workpiece_crate(self):
        """ Test get_workpiece_crate() method """
        wi = WarehouseInventory()
        color = WorkpieceColor.RED
        crate = Crate(Workpiece(0, color))
        crate_location = wi.insert(crate)
        loc = wi.get_loc_workpiece_crate(color)
        self.assertEqual(crate_location, loc)
        self.assertEqual(crate, wi.get(loc))

    def test_is_full(self):
        """ Test is_full() method """
        wi = WarehouseInventory()
        for _ in range(9):
            self.assertFalse(wi.is_full())
            wi.take(wi.get_loc_empty_crate())
            wi.insert(Crate(Workpiece(0, WorkpieceColor.RED)))
        self.assertTrue(wi.is_full())


if __name__ == '__main__':
    unittest.main()