start on both

hjwp · hjwp · commit f815b0d5733e · 2019-02-05T10:15:06.000Z
diff --git a/domain_model.py b/domain_model.py
@@ -41,11 +41,10 @@ def to_list(fn):
 @to_list
 def allocate(order, stock, shipments):
     for line in order.lines:
-        if stock:
+        for stock_line in stock:
             yield Allocation(order.id, line.sku, line.quantity, shipment_id=None)
-        else:
-            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)
+        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)
 
diff --git a/test_allocation.py b/test_allocation.py
@@ -67,8 +67,6 @@ def test_can_allocate_to_correct_shipment():
     assert allocations[0].quantity == 10
 
 
-
-
 def test_allocates_to_warehouse_stock_in_preference_to_shipment():
     sku = random_id()
     order = Order(id=random_id(), lines=[
@@ -118,3 +116,23 @@ def test_can_allocate_multiple_lines_to_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
 
+
+def test_can_allocate_to_both():
+    sku1, sku2 = random_id(), random_id()
+    order = Order(id=random_id(), lines=[
+        OrderLine(sku=sku1, quantity=10),
+        OrderLine(sku=sku2, quantity=10),
+    ])
+    shipment = Shipment(id=random_id(), eta=date.today(), lines=[
+        OrderLine(sku=sku2, quantity=1000),
+    ])
+    stock = [
+        Stock(sku=sku1, quantity=1000),
+    ]
+
+    allocations = allocate(order, stock, shipments=[shipment])
+    assert Allocation(order.id, sku1, 10, shipment_id=None) in allocations
+    assert Allocation(order.id, sku2, 10, shipment_id=shipment.id) in allocations
+
+
+
