new reqt re allocating to all if poss

hjwp · hjwp · commit ec48906ec3d8 · 2019-02-05T10:15:06.000Z
diff --git a/domain_model.py b/domain_model.py
@@ -19,6 +19,14 @@ class Shipment:
     lines: list
 
 
+def skus(thing):
+    try:
+        return {line.sku for line in thing}
+    except TypeError:
+        return {line.sku for line in thing.lines}
+
+
+
 def allocate_to_stock(line, stock):
     for stock_line in stock:
         if stock_line.sku == line.sku:
@@ -31,8 +39,18 @@ def allocate_to_shipments(line, shipments):
                 line.allocation = shipment.id
 
 def allocate(order, stock, shipments):
+    if skus(order) <= skus(stock):
+        for line in order:
+            line.allocation = 'STOCK'
+        return
+
+    for shipment in shipments:
+        if skus(order) <= skus(shipment):
+            for line in order:
+                line.allocation = shipment.id
+            return
+
     for line in order:
         allocate_to_stock(line, stock)
         if line.allocation is None:
             allocate_to_shipments(line, shipments)
-
diff --git a/test_allocation.py b/test_allocation.py
@@ -154,3 +154,21 @@ def test_can_allocate_to_both_preferring_stock():
     assert order[2].allocation == 'STOCK'
     assert order[3].allocation == 'STOCK'
 
+
+def test_mixed_allocations_are_avoided_if_possible():
+    sku1, sku2 = random_id(), random_id()
+    order = [
+        OrderLine(sku=sku1, quantity=10),
+        OrderLine(sku=sku2, quantity=10),
+    ]
+    shipment = Shipment(id=random_id(), eta=date.today(), lines=[
+        Line(sku=sku1, quantity=1000),
+        Line(sku=sku2, quantity=1000),
+    ])
+    stock = [Line(sku=sku1, quantity=1000)]
+
+    allocate(order, stock, shipments=[shipment])
+
+    assert order[0].allocation == shipment.id
+    assert order[1].allocation == shipment.id
+
