Add inventory functionality

elite_engineering/materials.py

@@ -1,3 +1,4 @@
+import collections
 import dataclasses
 import enum
 import fractions
@@ -153,8 +154,6 @@ class Material:
     category: str
     grade: int
 
-    MAX_CAPACITY: int = 300
-
     def trade_ratio(self, other: "Material") -> typing.Optional[fractions.Fraction]:
         if self.type_ != other.type_:
             return None
@@ -167,3 +166,18 @@ materials: typing.Dict[str, Material] = {
     for category, names in categories.items()
     for grade, name in enumerate(names)
 }
+
+
+class Inventory(collections.UserDict):
+    def __init__(self, items: typing.Optional[typing.Dict[str, int]] = None) -> None:
+        self.data = {name: 0 for name in materials}
+        if items:
+            self.update(items)
+
+    def __setitem__(self, key, value) -> None:
+        if key not in self.data:
+            raise KeyError(key)
+        super().__setitem__(key, value)
+
+    def has(self, items: typing.Dict[str, int]) -> bool:
+        return all([self[name] >= quantity for name, quantity in items.items()])

tests/test_materials.py

@@ -91,3 +91,35 @@ class TradeCalculationTests(unittest.TestCase):
         ]
 
         self.assertEqual(expected, actual)
+
+
+class InventoryTests(unittest.TestCase):
+    def test_create_empty_inventory(self) -> None:
+        inventory = materials.Inventory()
+        self.assertEqual(set(inventory.keys()), set(materials.materials.keys()))
+
+    def test_create_populated_inventory(self) -> None:
+        inventory = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+        self.assertEqual(5, inventory["Carbon"])
+        self.assertEqual(3, inventory["Shield Emitters"])
+        self.assertEqual(8, sum(inventory.values()))
+
+    def test_create_inventory_with_invalid_materials(self) -> None:
+        with self.assertRaises(KeyError) as exc:
+            materials.Inventory({"Carbon": 5, "Pizza": 12})
+        self.assertEqual(exc.exception.args, ("Pizza",))
+
+    def test_inventory_has_materials(self) -> None:
+        inventory = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+        self.assertTrue(inventory.has({"Carbon": 1, "Shield Emitters": 3}))
+
+    def test_inventory_lacks_materials(self) -> None:
+        inventory = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+        self.assertTrue(not inventory.has({"Nickel": 1}))
+        self.assertTrue(not inventory.has({"Carbon": 6, "Shield Emitters": 3}))
+
+    def test_inventory_has_raises_on_invalid_materials(self) -> None:
+        inventory = materials.Inventory()
+        with self.assertRaises(KeyError) as exc:
+            inventory.has({"Carbon": 5, "Pizza": 12})
+        self.assertEqual(exc.exception.args, ("Pizza",))