get_by_batchref in abstract+real repo, and a new integration test for it [get_by_batchref]

hjwp · hjwp · commit 6b2e8c6ee3e1 · 2019-11-20T21:12:36.000Z
diff --git a/src/allocation/repository.py b/src/allocation/repository.py
@@ -1,6 +1,7 @@
 from typing import Set
 import abc
-from allocation import model
+from allocation import model, orm
+
 
 
 class AbstractRepository(abc.ABC):
@@ -18,6 +19,12 @@ def get(self, sku) -> model.Product:
             self.seen.add(product)
         return product
 
+    def get_by_batchref(self, batchref) -> model.Product:
+        product = self._get_by_batchref(batchref)
+        if product:
+            self.seen.add(product)
+        return product
+
     @abc.abstractmethod
     def _add(self, product: model.Product):
         raise NotImplementedError
@@ -26,6 +33,11 @@ def _add(self, product: model.Product):
     def _get(self, sku) -> model.Product:
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def _get_by_batchref(self, batchref) -> model.Product:
+        raise NotImplementedError
+
+
 
 
 class SqlAlchemyRepository(AbstractRepository):
@@ -39,3 +51,8 @@ def _add(self, product):
 
     def _get(self, sku):
         return self.session.query(model.Product).filter_by(sku=sku).first()
+
+    def _get_by_batchref(self, batchref):
+        return self.session.query(model.Product).join(model.Batch).filter(
+            orm.batches.c.reference == batchref,
+        ).first()
diff --git a/tests/integration/test_repository.py b/tests/integration/test_repository.py
@@ -0,0 +1,14 @@
+from allocation import model, repository
+
+def test_get_by_batchref(session):
+    repo = repository.SqlAlchemyRepository(session)
+    b1 = model.Batch(ref='b1', sku='sku1', qty=100, eta=None)
+    b2 = model.Batch(ref='b2', sku='sku1', qty=100, eta=None)
+    b3 = model.Batch(ref='b3', sku='sku2', qty=100, eta=None)
+    p1 = model.Product(sku='sku1', batches=[b1, b2])
+    p2 = model.Product(sku='sku2', batches=[b3])
+    repo.add(p1)
+    repo.add(p2)
+    assert repo.get_by_batchref('b2') == p1
+    assert repo.get_by_batchref('b3') == p2
+
