classy version now at same stage

hjwp · hjwp · commit defc3e3573e7 · 2019-02-05T10:15:06.000Z
diff --git a/domain_model.py b/domain_model.py
@@ -14,13 +14,17 @@ def fully_allocated(self):
     def allocate(self, stock, shipments):
         self.allocations = {}
         for source in [stock] + shipments:
-            self.allocations.update({
-                key: stock
+            allocation = {
+                key: source
                 for key in self.lines
                 if key in source.lines
-            })
-            if self.fully_allocated:
+                and source.lines[key] > self.lines[key]
+            }
+            if set(allocation.keys()) == set(self.lines):
+                self.allocations = allocation
                 return
+            allocation.update(self.allocations)
+            self.allocations = allocation
 
 
 @dataclass
diff --git a/test_allocation.py b/test_allocation.py
@@ -130,15 +130,14 @@ def test_still_preserves_ordering_if_split_across_shipments():
     assert order.allocations['sku2'] == shipment2
     assert order.allocations['sku3'] == shipment2
 
-'''
 
 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 = Stock({
         'sku1': 1000,
         'sku2': 1000,
-    }
+    })
 
     order.allocate(stock, shipments=[shipment])
 
@@ -157,12 +156,11 @@ 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({)
+    shipment = Stock({
         'a-sku': 5,
-    }
+    })
 
     order.allocate(stock=Stock({}), shipments=[shipment])
 
     assert 'a-sku' not in order.allocations
 
-'''
