start on a classy version

hjwp · hjwp · commit b089c1b0b369 · 2019-02-05T10:15:06.000Z
diff --git a/domain_model.py b/domain_model.py
@@ -1,35 +1,26 @@
-def skus(d):
-    return set(d.keys())
+from dataclasses import dataclass
+from datetime import date
+
+@dataclass
+class Order:
+    lines: dict
+    allocations: dict = None
+
+    def allocate(self, stock, shipments):
+        self.allocations = {
+            key: stock
+            for key in self.lines
+        }
+        for s in shipments:
+            self.allocations.update({
+                key: s
+                for key in self.lines
+            })
+
+
+@dataclass
+class Stock:
+    lines: dict
+    eta: date = None
 
-def allocated_completely(order, allocation):
-    return skus(order) == skus(allocation)
-
-
-def allocate_to(order, source):
-    return {
-        sku: source
-        for sku, quantity in order.items()
-        if sku in source
-        and source[sku] > quantity
-    }
-
-
-def merge(allocations):
-    return {
-        k: v
-        for d in allocations
-        for k, v in d.items()
-    }
-
-
-
-def allocate(order, stock, shipments):
-    allocations = []
-    for source in [stock] + shipments:
-        allocation = allocate_to(order, source)
-        if allocated_completely(order, allocation):
-            return allocation
-        allocations.append(allocation)
-
-    return merge(reversed(allocations))
 
diff --git a/test_allocation.py b/test_allocation.py
@@ -1,33 +1,35 @@
-from domain_model import allocate
+from domain_model import Order, Stock
 
 
 def test_can_allocate_to_stock():
-    order = {'a-sku': 10}
-    stock = {'a-sku': 1000}
+    order = Order({'a-sku': 10})
+    stock = Stock({'a-sku': 1000})
 
-    allocations = allocate(order, stock, shipments=[])
+    order.allocate(stock, shipments=[])
 
-    assert allocations['a-sku'] == stock
+    assert order.allocations['a-sku'] == stock
 
 
 def test_can_allocate_to_shipment():
-    order = {'a-sku': 10}
-    shipment = {'a-sku': 1000}
+    order = Order({'a-sku': 10})
+    shipment = Stock({'a-sku': 1000})
 
-    allocations = allocate(order, stock={}, shipments=[shipment])
+    order.allocate(stock={}, shipments=[shipment])
 
-    assert allocations['a-sku'] == shipment
+    assert order.allocations['a-sku'] == shipment
 
 
 def test_ignores_irrelevant_stock():
-    order = {'sku1': 10}
-    stock = {'sku2': 1000}
-    shipment = {'sku1': 1000}
+    order = Order({'sku1': 10})
+    stock = Stock({'sku2': 1000})
+    shipment = Stock({'sku1': 1000})
 
-    allocations = allocate(order, stock=stock, shipments=[shipment])
+    order.allocate(stock=stock, shipments=[shipment])
+
+    assert order.allocations['sku1'] == shipment
 
-    assert allocations['sku1'] == shipment
 
+'''
 
 def test_can_allocate_to_correct_shipment():
     order = {'sku2': 10}
@@ -162,3 +164,4 @@ def test_cannot_allocate_if_insufficent_quantity_in_shipment():
 
     assert 'a-sku' not in allocations
 
+'''
