more like it, stock now allocates correctly

hjwp · hjwp · commit 16b01ec4e13d · 2019-02-05T10:15:06.000Z
diff --git a/domain_model.py b/domain_model.py
@@ -38,13 +38,24 @@ def to_list(fn):
     return lambda *a, **kw: list(fn(*a, **kw))
 
 
+def allocate_to_stock(order_id, line, stock):
+    for stock_line in stock:
+        if stock_line.sku == line.sku:
+            return Allocation(order_id, line.sku, line.quantity, shipment_id=None)
+
+def allocate_to_shipments(order_id, line, shipments):
+    for shipment in shipments:
+        for shipment_line in shipment.lines:
+            if shipment_line.sku == line.sku:
+                return Allocation(order_id, line.sku, line.quantity, shipment.id)
+
 @to_list
 def allocate(order, stock, shipments):
     for line in order.lines:
-        for stock_line in stock:
-            yield Allocation(order.id, line.sku, line.quantity, shipment_id=None)
-        for shipment in shipments:
-            for shipment_line in shipment.lines:
-                if shipment_line.sku == line.sku:
-                    yield Allocation(order.id, line.sku, line.quantity, shipment.id)
+        a = allocate_to_stock(order.id, line, stock)
+        if a:
+            yield a
+        else:
+            yield allocate_to_shipments(order.id, line, shipments)
+
 
diff --git a/test_allocation.py b/test_allocation.py
@@ -46,6 +46,23 @@ def test_can_allocate_to_shipment():
     assert allocations[0].shipment_id == shipment.id
     assert allocations[0].quantity == 10
 
+def test_ignores_invalid_stock():
+    sku1, sku2 = random_id(), random_id()
+    order = Order(id=random_id(), lines=[
+        OrderLine(sku=sku1, quantity=10),
+    ])
+    stock = Stock(sku=sku2, quantity=1000)
+    shipment = Shipment(id=random_id(), eta=date.today(), lines=[
+        OrderLine(sku=sku1, quantity=1000),
+    ])
+
+    allocations = allocate(order, stock=[stock], shipments=[shipment])
+
+    assert allocations[0].order_id == order.id
+    assert allocations[0].sku == sku1
+    assert allocations[0].shipment_id == shipment.id
+    assert allocations[0].quantity == 10
+
 
 def test_can_allocate_to_correct_shipment():
     sku1, sku2 = random_id(), random_id()
@@ -135,4 +152,30 @@ def test_can_allocate_to_both():
     assert Allocation(order.id, sku2, 10, shipment_id=shipment.id) in allocations
 
 
+def test_can_allocate_to_both_preferring_stock():
+    sku1, sku2, sku3, sku4 = [random_id() for _ in range(4)]
+    order = Order(id=random_id(), lines=[
+        OrderLine(sku=sku1, quantity=10),
+        OrderLine(sku=sku2, quantity=10),
+        OrderLine(sku=sku3, quantity=10),
+        OrderLine(sku=sku4, quantity=10),
+    ])
+    shipment = Shipment(id=random_id(), eta=date.today(), lines=[
+        OrderLine(sku=sku1, quantity=1000),
+        OrderLine(sku=sku2, quantity=1000),
+        OrderLine(sku=sku3, quantity=1000),
+    ])
+    stock = [
+        Stock(sku=sku3, quantity=1000),
+        Stock(sku=sku4, quantity=1000),
+    ]
+
+    allocations = allocate(order, stock, shipments=[shipment])
+    assert Allocation(order.id, sku1, 10, shipment_id=shipment.id) in allocations
+    assert Allocation(order.id, sku2, 10, shipment_id=shipment.id) in allocations
+    assert Allocation(order.id, sku3, 10, shipment_id=None) in allocations
+    assert Allocation(order.id, sku4, 10, shipment_id=None) in allocations
+    assert Allocation(order.id, sku3, 10, shipment_id=shipment.id) not in allocations
+
+
 
