ORM for _allocations set on Batch

hjwp · hjwp · commit fff68bd217c0 · 2021-02-23T16:37:24.000Z
diff --git a/orm.py b/orm.py
@@ -1,5 +1,5 @@
-from sqlalchemy import Table, MetaData, Column, Integer, String, Date
-from sqlalchemy.orm import mapper
+from sqlalchemy import Table, MetaData, Column, Integer, String, Date, ForeignKey
+from sqlalchemy.orm import mapper, relationship
 
 import model
 
@@ -9,21 +9,39 @@
 order_lines = Table(
     "order_lines",
     metadata,
-    Column("orderid", String(255), primary_key=True),
-    Column("sku", String(255), primary_key=True),
-    Column("qty", Integer),
+    Column("id", Integer, primary_key=True, autoincrement=True),
+    Column("sku", String(255)),
+    Column("qty", Integer, nullable=False),
+    Column("orderid", String(255)),
 )
 
 batches = Table(
     "batches",
     metadata,
-    Column("reference", String(255), primary_key=True),
-    Column("sku", String(255), primary_key=True),
-    Column("_purchased_qty", Integer),
-    Column("eta", Date),
+    Column("id", Integer, primary_key=True, autoincrement=True),
+    Column("reference", String(255)),
+    Column("sku", String(255)),
+    Column("_purchased_quantity", Integer, nullable=False),
+    Column("eta", Date, nullable=True),
+)
+
+allocations = Table(
+    "allocations",
+    metadata,
+    Column("id", Integer, primary_key=True, autoincrement=True),
+    Column("orderline_id", ForeignKey("order_lines.id")),
+    Column("batch_id", ForeignKey("batches.id")),
 )
 
 
 def start_mappers():
-    mapper(model.OrderLine, order_lines)
-    mapper(model.Batch, batches)
+    lines_mapper = mapper(model.OrderLine, order_lines)
+    mapper(
+        model.Batch,
+        batches,
+        properties={
+            "_allocations": relationship(
+                lines_mapper, secondary=allocations, collection_class=set,
+            )
+        },
+    )
diff --git a/test_orm.py b/test_orm.py
@@ -26,14 +26,66 @@ def test_orderline_mapper_can_save_lines(session):
     assert rows == [("order1", "DECORATIVE-WIDGET", 12)]
 
 
-def test_batches(session):
-    session.execute('INSERT INTO "batches" VALUES ("batch1", "sku1", 100, null)')
+def test_retrieving_batches(session):
     session.execute(
-        'INSERT INTO "batches" VALUES ("batch2", "sku2", 200, "2011-04-11")'
+        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+        ' VALUES ("batch1", "sku1", 100, null)'
+    )
+    session.execute(
+        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+        ' VALUES ("batch2", "sku2", 200, "2011-04-11")'
     )
     expected = [
         model.Batch("batch1", "sku1", 100, eta=None),
         model.Batch("batch2", "sku2", 200, eta=date(2011, 4, 11)),
     ]
 
     assert session.query(model.Batch).all() == expected
+
+
+def test_saving_batches(session):
+    batch = model.Batch("batch1", "sku1", 100, eta=None)
+    session.add(batch)
+    session.commit()
+    rows = list(
+        session.execute(
+            'SELECT reference, sku, _purchased_quantity, eta FROM "batches"'
+        )
+    )
+    assert rows == [("batch1", "sku1", 100, None)]
+
+
+def test_saving_allocations(session):
+    batch = model.Batch("batch1", "sku1", 100, eta=None)
+    line = model.OrderLine("order1", "sku1", 10)
+    batch.allocate(line)
+    session.add(batch)
+    session.commit()
+    rows = list(session.execute('SELECT orderline_id, batch_id FROM "allocations"'))
+    assert rows == [(batch.id, line.id)]
+
+
+def test_retrieving_allocations(session):
+    session.execute(
+        'INSERT INTO order_lines (orderid, sku, qty) VALUES ("order1", "sku1", 12)'
+    )
+    [[olid]] = session.execute(
+        "SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku",
+        dict(orderid="order1", sku="sku1"),
+    )
+    session.execute(
+        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+        ' VALUES ("batch1", "sku1", 100, null)'
+    )
+    [[bid]] = session.execute(
+        "SELECT id FROM batches WHERE reference=:ref AND sku=:sku",
+        dict(ref="batch1", sku="sku1"),
+    )
+    session.execute(
+        "INSERT INTO allocations (orderline_id, batch_id) VALUES (:olid, :bid)",
+        dict(olid=olid, bid=bid),
+    )
+
+    batch = session.query(model.Batch).one()
+
+    assert batch._allocations == {model.OrderLine("order1", "sku1", 12)}
