update the UoW to work with django. quite a lot of work

hjwp · hjwp · commit 06e58834e0ba · 2021-02-24T20:23:43.000Z
diff --git a/src/allocation/adapters/repository.py b/src/allocation/adapters/repository.py
@@ -23,14 +23,19 @@ def _get(self, sku):
 
 class DjangoRepository(AbstractRepository):
     def __init__(self):
+        super().__init__()
         from djangoproject.alloc import models
 
         self.django_models = models
 
     def add(self, batch):
+        super().add(batch)
+        self.update(batch)
+
+    def update(self, batch):
         self.django_models.Batch.update_from_domain(batch)
 
-    def get(self, reference):
+    def _get(self, reference):
         return (
             self.django_models.Batch.objects.filter(reference=reference)
             .first()
diff --git a/src/allocation/service_layer/unit_of_work.py b/src/allocation/service_layer/unit_of_work.py
@@ -1,11 +1,7 @@
 # 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 django.db import transaction
 from allocation.adapters import repository
 
 
@@ -27,28 +23,20 @@ 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
-
+class DjangoUnitOfWork(AbstractUnitOfWork):
     def __enter__(self):
-        self.session = self.session_factory()  # type: Session
-        self.batches = repository.SqlAlchemyRepository(self.session)
+        self.batches = repository.DjangoRepository()
+        transaction.set_autocommit(False)
         return super().__enter__()
 
     def __exit__(self, *args):
         super().__exit__(*args)
-        self.session.close()
+        transaction.set_autocommit(True)
 
     def commit(self):
-        self.session.commit()
+        for batch in self.batches.seen:
+            self.batches.update(batch)
+        transaction.commit()
 
     def rollback(self):
-        self.session.rollback()
+        transaction.rollback()
diff --git a/src/djangoproject/alloc/models.py b/src/djangoproject/alloc/models.py
@@ -8,15 +8,16 @@ class Batch(models.Model):
     qty = models.IntegerField()
     eta = models.DateField(blank=True, null=True)
 
-    @classmethod
-    def update_from_domain(self, batch: domain_model.Batch):
+    @staticmethod
+    def update_from_domain(batch: domain_model.Batch):
         try:
             b = Batch.objects.get(reference=batch.reference)
         except Batch.DoesNotExist:
             b = Batch(reference=batch.reference)
         b.sku = batch.sku
         b.qty = batch._purchased_quantity
         b.eta = batch.eta
+        b.allocation_set.set(Allocation.from_domain(l, b) for l in batch._allocations)
         b.save()
 
     def to_domain(self) -> domain_model.Batch:
@@ -37,7 +38,22 @@ def to_domain(self):
             orderid=self.orderid, sku=self.sku, qty=self.qty
         )
 
+    @staticmethod
+    def from_domain(line):
+        l, _ = OrderLine.objects.get_or_create(
+            orderid=line.orderid, sku=line.sku, qty=line.qty
+        )
+        return l
+
 
 class Allocation(models.Model):
     batch = models.ForeignKey(Batch, on_delete=models.CASCADE)
     line = models.ForeignKey(OrderLine, on_delete=models.CASCADE)
+
+    @staticmethod
+    def from_domain(domain_line, django_batch):
+        a, _ = Allocation.objects.get_or_create(
+            line=OrderLine.from_domain(domain_line),
+            batch=django_batch,
+        )
+        return a
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -7,7 +7,13 @@
 from requests.exceptions import ConnectionError
 
 from allocation import config
-from allocation.adapters.orm import metadata, start_mappers
+
+
+@pytest.fixture
+def django_models():
+    from djangoproject.alloc import models
+
+    return models
 
 
 def wait_for_postgres_to_come_up(engine):
diff --git a/tests/integration/test_repository.py b/tests/integration/test_repository.py
@@ -5,13 +5,6 @@
 from allocation.adapters import repository
 
 
-@pytest.fixture
-def django_models():
-    from djangoproject.alloc import models
-
-    return models
-
-
 @pytest.mark.django_db
 def test_repository_can_save_a_batch(django_models):
     batch = model.Batch("batch1", "RUSTY-SOAPDISH", 100, eta=date(2011, 12, 25))
diff --git a/tests/integration/test_uow.py b/tests/integration/test_uow.py
@@ -3,63 +3,52 @@
 from allocation.service_layer 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 insert_batch(django_models, ref, sku, qty, eta):
+    django_models.Batch.objects.create(reference=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 get_allocated_batch_ref(django_models, orderid, sku):
+    print(django_models.Allocation.objects.all())
+    return django_models.Allocation.objects.get(
+        line__orderid=orderid, line__sku=sku
+    ).batch.reference
 
 
-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()
+@pytest.mark.django_db(transaction=True)
+def test_uow_can_retrieve_a_batch_and_allocate_to_it(django_models):
+    insert_batch(django_models, "batch1", "HIPSTER-WORKBENCH", 100, None)
 
-    uow = unit_of_work.SqlAlchemyUnitOfWork(session_factory)
+    uow = unit_of_work.DjangoUnitOfWork()
     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")
+    batchref = get_allocated_batch_ref(django_models, "o1", "HIPSTER-WORKBENCH")
     assert batchref == "batch1"
 
 
-def test_rolls_back_uncommitted_work_by_default(session_factory):
-    uow = unit_of_work.SqlAlchemyUnitOfWork(session_factory)
+@pytest.mark.django_db(transaction=True)
+def test_rolls_back_uncommitted_work_by_default(django_models):
+    uow = unit_of_work.DjangoUnitOfWork()
     with uow:
-        insert_batch(uow.session, "batch1", "MEDIUM-PLINTH", 100, None)
+        insert_batch(django_models, "batch1", "MEDIUM-PLINTH", 100, None)
 
-    new_session = session_factory()
-    rows = list(new_session.execute('SELECT * FROM "batches"'))
-    assert rows == []
+    rows = django_models.Batch.objects.all()
+    assert list(rows) == []
 
 
-def test_rolls_back_on_error(session_factory):
+@pytest.mark.django_db(transaction=True)
+def test_rolls_back_on_error(django_models):
     class MyException(Exception):
         pass
 
-    uow = unit_of_work.SqlAlchemyUnitOfWork(session_factory)
+    uow = unit_of_work.DjangoUnitOfWork()
     with pytest.raises(MyException):
         with uow:
-            insert_batch(uow.session, "batch1", "LARGE-FORK", 100, None)
+            insert_batch(django_models, "batch1", "LARGE-FORK", 100, None)
             raise MyException()
 
-    new_session = session_factory()
-    rows = list(new_session.execute('SELECT * FROM "batches"'))
-    assert rows == []
+    rows = django_models.Batch.objects.all()
+    assert list(rows) == []
