diff --git a/elite_engineering/materials.py b/elite_engineering/materials.py index 2955c7f..93ddb9e 100644 --- a/elite_engineering/materials.py +++ b/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()]) diff --git a/tests/test_materials.py b/tests/test_materials.py index af728f6..c2deb9a 100644 --- a/tests/test_materials.py +++ b/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",))