Move document creation into a class

Author: jerome-white

backend/app/tests/crud/documents/_utils.py

@@ -11,15 +11,11 @@
 from app.crud.user import get_user_by_email
 from app.models import Document, UserCreate
 
-class Constants:
-    n_documents = 10
-
 @ft.cache
 def get_user_id_by_email(session: Session):
     user = get_user_by_email(session=session, email=settings.FIRST_SUPERUSER)
     return user.id
 
-
 @ft.cache
 def int_to_uuid(value):
     return UUID(int=value)
@@ -28,31 +24,44 @@ def rm_documents(session: Session):
     session.exec(delete(Document))
     session.commit()
 
-def mk_document(owner_id, index=0):
-    doc_id = int_to_uuid(index)
-
-    args = str(doc_id).split('-')
-    fname = Path('/', *args).with_suffix('.xyz')
-    return Document(
-        id=doc_id,
-        owner_id=owner_id,
-        fname=fname.name,
-        object_store_url=fname.as_uri(),
-    )
-
 def insert_documents(session: Session, n: int):
-    owner_id = get_user_id_by_email(session)
-
     crud = DocumentCrud(session)
-    for i in range(n):
-        document = mk_document(owner_id, i)
+    docs = DocumentMaker(session)
 
-        session.add(document)
+    for (_, d) in zip(range(n), docs):
+        session.add(d)
         session.commit()
-        session.refresh(document)
-
-        yield document
+        session.refresh(d)
+        yield d
 
 def insert_document(session: Session):
     (document, ) = insert_documents(session, 1)
     return document
+
+class Constants:
+    n_documents = 10
+
+class DocumentMaker:
+    def __init__(self, session: Session):
+        self.owner_id = get_user_id_by_email(session)
+        self.index = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        doc_id = self.get_and_increment()
+        args = str(doc_id).split('-')
+        fname = Path('/', *args).with_suffix('.xyz')
+
+        return Document(
+            id=doc_id,
+            owner_id=self.owner_id,
+            fname=fname.name,
+            object_store_url=fname.as_uri(),
+        )
+
+    def get_and_increment(self):
+        doc_id = int_to_uuid(self.index)
+        self.index += 1
+        return doc_id

backend/app/tests/crud/documents/test_update.py

@@ -1,53 +1,44 @@
 from uuid import UUID
 from typing import ClassVar
-from dataclasses import dataclass
 
 import pytest
 from sqlmodel import Session
 
 from app.crud import DocumentCrud
 
 from _utils import (
-    get_user_id_by_email,
-    mk_document,
+    DocumentMaker,
     rm_documents,
 )
 
-@dataclass
-class State:
-    crud: DocumentCrud
-    owner_id: UUID
-    _doc_id: ClassVar[int] = 0
+class TestState:
+    def __init__(self, db: Session):
+        self.crud = DocumentCrud(db)
+        self.documents = DocumentMaker(db)
 
     def add(self):
-        document = mk_document(self.owner_id, self._doc_id)
-        self._doc_id += 1
-        return self.crud.update(document)
+        return self.crud.update(next(self.documents))
 
     def get(self):
-        return self.crud.read_many(self.owner_id)
+        return self.crud.read_many(self.documents.owner_id)
 
 @pytest.fixture
 def state(db: Session):
     rm_documents(db)
-
-    crud = DocumentCrud(db)
-    owner_id = get_user_id_by_email(db)
-
-    return State(crud, owner_id)
+    return TestState(db)
 
 class TestDatabaseUpdate:
-    def test_update_adds_one(self, state: State):
+    def test_update_adds_one(self, state: TestState):
         before = state.get()
         state.add()
         after = state.get()
 
         assert len(before) + 1 == len(after)
 
-    def test_sequential_update_is_ordered(self, state: State):
+    def test_sequential_update_is_ordered(self, state: TestState):
         (a, b) = (state.add() for _ in range(2))
         assert a.created_at <= b.created_at
 
-    def test_insert_does_not_delete(self, state: State):
+    def test_insert_does_not_delete(self, state: TestState):
         document = state.add()
         assert document.deleted_at is None