basic uow test, uow and conftest.py changes

Author: hjwp

src/allocation/service_layer/services.py

@@ -16,19 +16,16 @@ def is_valid_sku(sku, batches):
 
 
 def add_batch(
-    ref: str,
-    sku: str,
-    qty: int,
-    eta: Optional[date],
-    repo: AbstractRepository,
-    session,
+    ref: str, sku: str, qty: int, eta: Optional[date],
+    repo: AbstractRepository, session,
 ) -> None:
     repo.add(model.Batch(ref, sku, qty, eta))
     session.commit()
 
 
 def allocate(
-    orderid: str, sku: str, qty: int, repo: AbstractRepository, session
+    orderid: str, sku: str, qty: int,
+    repo: AbstractRepository, session,
 ) -> str:
     line = OrderLine(orderid, sku, qty)
     batches = repo.list()

src/allocation/unit_of_work.py

@@ -0,0 +1,50 @@
+# pylint: disable=attribute-defined-outside-init
+from __future__ import annotations
+import abc
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm.session import Session
+
+from allocation import config
+from allocation import repository
+
+
+class AbstractUnitOfWork(abc.ABC):
+    batches: repository.AbstractRepository
+
+    def __enter__(self) -> AbstractUnitOfWork:
+        return self
+
+    def __exit__(self, *args):
+        self.rollback()
+
+    @abc.abstractmethod
+    def commit(self):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def rollback(self):
+        raise NotImplementedError
+
+
+DEFAULT_SESSION_FACTORY = sessionmaker(bind=create_engine(config.get_postgres_uri(),))
+
+
+class SqlAlchemyUnitOfWork(AbstractUnitOfWork):
+    def __init__(self, session_factory=DEFAULT_SESSION_FACTORY):
+        self.session_factory = session_factory
+
+    def __enter__(self):
+        self.session = self.session_factory()  # type: Session
+        self.batches = repository.SqlAlchemyRepository(self.session)
+        return super().__enter__()
+
+    def __exit__(self, *args):
+        super().__exit__(*args)
+        self.session.close()
+
+    def commit(self):
+        self.session.commit()
+
+    def rollback(self):
+        self.session.rollback()

tests/conftest.py

@@ -21,12 +21,17 @@ def in_memory_db():
 
 
 @pytest.fixture
-def session(in_memory_db):
+def session_factory(in_memory_db):
     start_mappers()
-    yield sessionmaker(bind=in_memory_db)()
+    yield sessionmaker(bind=in_memory_db)
     clear_mappers()
 
 
+@pytest.fixture
+def session(session_factory):
+    return session_factory()
+
+
 def wait_for_postgres_to_come_up(engine):
     deadline = time.time() + 10
     while time.time() < deadline:

tests/integration/test_uow.py

@@ -0,0 +1,40 @@
+import pytest
+from allocation import model
+from allocation import unit_of_work
+
+
+def insert_batch(session, ref, sku, qty, eta):
+    session.execute(
+        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+        " VALUES (:ref, :sku, :qty, :eta)",
+        dict(ref=ref, sku=sku, qty=qty, eta=eta),
+    )
+
+
+def get_allocated_batch_ref(session, orderid, sku):
+    [[orderlineid]] = session.execute(
+        "SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku",
+        dict(orderid=orderid, sku=sku),
+    )
+    [[batchref]] = session.execute(
+        "SELECT b.reference FROM allocations JOIN batches AS b ON batch_id = b.id"
+        " WHERE orderline_id=:orderlineid",
+        dict(orderlineid=orderlineid),
+    )
+    return batchref
+
+
+def test_uow_can_retrieve_a_batch_and_allocate_to_it(session_factory):
+    session = session_factory()
+    insert_batch(session, "batch1", "HIPSTER-WORKBENCH", 100, None)
+    session.commit()
+
+    uow = unit_of_work.SqlAlchemyUnitOfWork(session_factory)
+    with uow:
+        batch = uow.batches.get(reference="batch1")
+        line = model.OrderLine("o1", "HIPSTER-WORKBENCH", 10)
+        batch.allocate(line)
+        uow.commit()
+
+    batchref = get_allocated_batch_ref(session, "o1", "HIPSTER-WORKBENCH")
+    assert batchref == "batch1"