implement .seen on repository [repository_tracks_seen]

hjwp · hjwp · commit 803d9b69c1e6 · 2019-11-20T21:12:36.000Z
diff --git a/src/allocation/repository.py b/src/allocation/repository.py
@@ -1,32 +1,41 @@
-import abc
 from typing import Set
+import abc
 from allocation import model
 
 
 class AbstractRepository(abc.ABC):
 
-    @abc.abstractmethod
+    def __init__(self):
+        self.seen = set()  # type: Set[model.Product]
+
     def add(self, product: model.Product):
+        self._add(product)
+        self.seen.add(product)
+
+    def get(self, sku) -> model.Product:
+        product = self._get(sku)
+        if product:
+            self.seen.add(product)
+        return product
+
+    @abc.abstractmethod
+    def _add(self, product: model.Product):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def get(self, sku) -> model.Product:
+    def _get(self, sku) -> model.Product:
         raise NotImplementedError
 
 
 
 class SqlAlchemyRepository(AbstractRepository):
 
     def __init__(self, session):
+        super().__init__()
         self.session = session
-        self.seen = set()  # type: Set[model.Product]
 
-    def add(self, product):
-        self.seen.add(product)
+    def _add(self, product):
         self.session.add(product)
 
-    def get(self, sku):
-        product = self.session.query(model.Product).filter_by(sku=sku).first()
-        if product:
-            self.seen.add(product)
-        return product
+    def _get(self, sku):
+        return self.session.query(model.Product).filter_by(sku=sku).first()
diff --git a/tests/unit/test_services.py b/tests/unit/test_services.py
@@ -6,12 +6,13 @@
 class FakeRepository(repository.AbstractRepository):
 
     def __init__(self, products):
+        super().__init__()
         self._products = set(products)
 
-    def add(self, product):
+    def _add(self, product):
         self._products.add(product)
 
-    def get(self, sku):
+    def _get(self, sku):
         return next((p for p in self._products if p.sku == sku), None)
 
 
