separate out a shipment class, remove lines attr

hjwp · hjwp · commit dc06fe1d106c · 2019-02-05T10:54:18.000Z
diff --git a/domain_model.py b/domain_model.py
@@ -37,11 +37,19 @@ def allocate(self, stock, shipments):
             self.allocation = allocation
 
 
+class Stock(dict):
+
+    def can_allocate(self, sku, quantity):
+        return sku in self and self[sku] > quantity
+
+
 @dataclass
-class Stock:
-    lines: dict
-    eta: date = None
+class Shipment(Stock):
+
+    def __init__(self, lines, eta):
+        self.eta = eta
+        super().__init__(lines)
 
     def can_allocate(self, sku, quantity):
-        return sku in self.lines and self.lines[sku] > quantity
+        return sku in self and self[sku] > quantity
 
diff --git a/test_allocation.py b/test_allocation.py
@@ -1,4 +1,4 @@
-from domain_model import Order, Stock
+from domain_model import Order, Stock, Shipment
 from datetime import date, timedelta
 
 today = date.today()
@@ -17,7 +17,7 @@ def test_can_allocate_to_stock():
 
 def test_can_allocate_to_shipment():
     order = Order({'a-sku': 10})
-    shipment = Stock({'a-sku': 1000}, eta=tomorrow)
+    shipment = Shipment({'a-sku': 1000}, eta=tomorrow)
 
     order.allocate(stock=Stock({}), shipments=[shipment])
 
@@ -27,7 +27,7 @@ def test_can_allocate_to_shipment():
 def test_ignores_irrelevant_stock():
     order = Order({'sku1': 10})
     stock = Stock({'sku2': 1000})
-    shipment = Stock({'sku1': 1000}, eta=tomorrow)
+    shipment = Shipment({'sku1': 1000}, eta=tomorrow)
 
     order.allocate(stock=stock, shipments=[shipment])
 
@@ -37,8 +37,8 @@ def test_ignores_irrelevant_stock():
 
 def test_can_allocate_to_correct_shipment():
     order = Order({'sku2': 10})
-    shipment1 = Stock({'sku1': 1000}, eta=tomorrow)
-    shipment2 = Stock({'sku2': 1000}, eta=tomorrow)
+    shipment1 = Shipment({'sku1': 1000}, eta=tomorrow)
+    shipment2 = Shipment({'sku2': 1000}, eta=tomorrow)
 
     order.allocate(stock=Stock({}), shipments=[shipment1, shipment2])
 
@@ -48,7 +48,7 @@ def test_can_allocate_to_correct_shipment():
 def test_allocates_to_stock_in_preference_to_shipment():
     order = Order({'sku1': 10})
     stock = Stock({'sku1': 1000})
-    shipment = Stock({'sku1': 1000}, eta=tomorrow)
+    shipment = Shipment({'sku1': 1000}, eta=tomorrow)
 
     order.allocate(stock, shipments=[shipment])
 
@@ -66,7 +66,7 @@ def test_can_allocate_multiple_lines_to_wh():
 
 def test_can_allocate_multiple_lines_to_shipment():
     order = Order({'sku1': 5, 'sku2': 10})
-    shipment = Stock({'sku1': 1000, 'sku2': 1000}, eta=tomorrow)
+    shipment = Shipment({'sku1': 1000, 'sku2': 1000}, eta=tomorrow)
 
     order.allocate(stock=Stock({}), shipments=[shipment])
 
@@ -76,7 +76,7 @@ def test_can_allocate_multiple_lines_to_shipment():
 
 def test_can_allocate_to_both():
     order = Order({'sku1': 5, 'sku2': 10})
-    shipment = Stock({'sku2': 1000}, eta=tomorrow)
+    shipment = Shipment({'sku2': 1000}, eta=tomorrow)
     stock = Stock({'sku1': 1000})
 
     order.allocate(stock, shipments=[shipment])
@@ -87,7 +87,7 @@ def test_can_allocate_to_both():
 
 def test_can_allocate_to_both_preferring_stock():
     order = Order({'sku1': 1, 'sku2': 2, 'sku3': 3, 'sku4': 4})
-    shipment = Stock({'sku1': 1000, 'sku2': 1000, 'sku3': 1000}, eta=tomorrow)
+    shipment = Shipment({'sku1': 1000, 'sku2': 1000, 'sku3': 1000}, eta=tomorrow)
     stock = Stock({'sku3': 1000, 'sku4': 1000})
 
     order.allocate(stock, shipments=[shipment])
@@ -100,7 +100,7 @@ def test_can_allocate_to_both_preferring_stock():
 
 def test_mixed_allocation_are_avoided_if_possible():
     order = Order({'sku1': 10, 'sku2': 10})
-    shipment = Stock({'sku1': 1000, 'sku2': 1000}, eta=tomorrow)
+    shipment = Shipment({'sku1': 1000, 'sku2': 1000}, eta=tomorrow)
     stock = Stock({'sku1': 1000})
 
     order.allocate(stock, shipments=[shipment])
@@ -111,8 +111,8 @@ def test_mixed_allocation_are_avoided_if_possible():
 
 def test_allocated_to_earliest_suitable_shipment_in_list():
     order = Order({'sku1': 10, 'sku2': 10})
-    shipment1 = Stock({'sku1': 1000, 'sku2': 1000}, eta=today)
-    shipment2 = Stock({'sku1': 1000, 'sku2': 1000}, eta=tomorrow)
+    shipment1 = Shipment({'sku1': 1000, 'sku2': 1000}, eta=today)
+    shipment2 = Shipment({'sku1': 1000, 'sku2': 1000}, eta=tomorrow)
     stock = Stock({})
 
     order.allocate(stock, shipments=[shipment2, shipment1])
@@ -123,9 +123,9 @@ def test_allocated_to_earliest_suitable_shipment_in_list():
 
 def test_still_chooses_earliest_if_split_across_shipments():
     order = Order({'sku1': 10, 'sku2': 10, 'sku3': 10})
-    shipment1 = Stock({'sku1': 1000}, eta=today)
-    shipment2 = Stock({'sku2': 1000, 'sku3': 1000}, eta=tomorrow)
-    shipment3 = Stock({'sku2': 1000, 'sku3': 1000}, eta=later)
+    shipment1 = Shipment({'sku1': 1000}, eta=today)
+    shipment2 = Shipment({'sku2': 1000, 'sku3': 1000}, eta=tomorrow)
+    shipment3 = Shipment({'sku2': 1000, 'sku3': 1000}, eta=later)
     stock = Stock({})
 
     order.allocate(stock, shipments=[shipment3, shipment2, shipment1])
@@ -138,10 +138,10 @@ def test_still_chooses_earliest_if_split_across_shipments():
 def test_stock_not_quite_enough_means_we_use_shipment():
     order = Order({'sku1': 10, 'sku2': 10})
     stock = Stock({'sku1': 10, 'sku2': 5})
-    shipment = Stock({
+    shipment = Shipment({
         'sku1': 1000,
         'sku2': 1000,
-    })
+    }, eta=tomorrow)
 
     order.allocate(stock, shipments=[shipment])
 
@@ -160,7 +160,7 @@ def test_cannot_allocate_if_insufficent_quantity_in_stock():
 
 def test_cannot_allocate_if_insufficent_quantity_in_shipment():
     order = Order({'a-sku': 10})
-    shipment = Stock({'a-sku': 5}, eta=tomorrow)
+    shipment = Shipment({'a-sku': 5}, eta=tomorrow)
 
     order.allocate(stock=Stock({}), shipments=[shipment])
 
