Refactor remaining test.utils helpers into actual tests.

[plomtask] / plomtask / db.py

diff --git a/plomtask/db.py b/plomtask/db.py
index 99998a6ab29f760ba0d62f90395739dad4b521ff..cce2630cd58bfb8bf7283c2eb0d2f45006d8ba26 100644 (file)
--- a/plomtask/db.py
+++ b/plomtask/db.py
@@ -238,6 +238,7 @@ class BaseModel(Generic[BaseModelId]):
     id_: None | BaseModelId
     cache_: dict[BaseModelId, Self]
     to_search: list[str] = []
+    can_create_by_id = False
     _exists = True
 
     def __init__(self, id_: BaseModelId | None) -> None:
@@ -271,6 +272,26 @@ class BaseModel(Generic[BaseModelId]):
         assert isinstance(other.id_, int)
         return self.id_ < other.id_
 
+    @property
+    def as_dict(self) -> dict[str, object]:
+        """Return self as (json.dumps-coompatible) dict."""
+        d: dict[str, object] = {'id': self.id_}
+        if len(self.to_save_versioned) > 0:
+            d['_versioned'] = {}
+        for k in self.to_save:
+            attr = getattr(self, k)
+            if hasattr(attr, 'as_dict'):
+                d[k] = attr.as_dict
+            d[k] = attr
+        for k in self.to_save_versioned:
+            attr = getattr(self, k)
+            assert isinstance(d['_versioned'], dict)
+            d['_versioned'][k] = attr.history
+        for r in self.to_save_relations:
+            attr_name = r[2]
+            d[attr_name] = [x.as_dict for x in getattr(self, attr_name)]
+        return d
+
     # cache management
     # (we primarily use the cache to ensure we work on the same object in
     # memory no matter where and how we retrieve it, e.g. we don't want
@@ -295,7 +316,13 @@ class BaseModel(Generic[BaseModelId]):
 
     @classmethod
     def empty_cache(cls) -> None:
-        """Empty class's cache."""
+        """Empty class's cache, and disappear all former inhabitants."""
+        # pylint: disable=protected-access
+        # (cause we remain within the class)
+        if hasattr(cls, 'cache_'):
+            to_disappear = list(cls.cache_.values())
+            for item in to_disappear:
+                item._disappear()
         cls.cache_ = {}
 
     @classmethod
@@ -310,15 +337,14 @@ class BaseModel(Generic[BaseModelId]):
     def _get_cached(cls: type[BaseModelInstance],
                     id_: BaseModelId) -> BaseModelInstance | None:
         """Get object of id_ from class's cache, or None if not found."""
-        # pylint: disable=consider-iterating-dictionary
         cache = cls.get_cache()
-        if id_ in cache.keys():
+        if id_ in cache:
             obj = cache[id_]
             assert isinstance(obj, cls)
             return obj
         return None
 
-    def _cache(self) -> None:
+    def cache(self) -> None:
         """Update object in class's cache.
 
         Also calls ._disappear if cache holds older reference to object of same
@@ -349,22 +375,23 @@ class BaseModel(Generic[BaseModelId]):
                        # pylint: disable=unused-argument
                        db_conn: DatabaseConnection,
                        row: Row | list[Any]) -> BaseModelInstance:
-        """Make from DB row, update DB cache with it."""
+        """Make from DB row (sans relations), update DB cache with it."""
         obj = cls(*row)
-        obj._cache()
+        assert obj.id_ is not None
+        for attr_name in cls.to_save_versioned:
+            attr = getattr(obj, attr_name)
+            table_name = attr.table_name
+            for row_ in db_conn.row_where(table_name, 'parent', obj.id_):
+                attr.history_from_row(row_)
+        obj.cache()
         return obj
 
     @classmethod
-    def by_id(cls, db_conn: DatabaseConnection,
-              id_: BaseModelId | None,
-              # pylint: disable=unused-argument
-              create: bool = False) -> Self:
+    def by_id(cls, db_conn: DatabaseConnection, id_: BaseModelId) -> Self:
         """Retrieve by id_, on failure throw NotFoundException.
 
         First try to get from cls.cache_, only then check DB; if found,
         put into cache.
-
-        If create=True, make anew (but do not cache yet).
         """
         obj = None
         if id_ is not None:
@@ -375,11 +402,22 @@ class BaseModel(Generic[BaseModelId]):
                     break
         if obj:
             return obj
-        if create:
-            obj = cls(id_)
-            return obj
         raise NotFoundException(f'found no object of ID {id_}')
 
+    @classmethod
+    def by_id_or_create(cls, db_conn: DatabaseConnection,
+                        id_: BaseModelId | None
+                        ) -> Self:
+        """Wrapper around .by_id, creating (not caching/saving) if not find."""
+        if not cls.can_create_by_id:
+            raise HandledException('Class cannot .by_id_or_create.')
+        if id_ is None:
+            return cls(None)
+        try:
+            return cls.by_id(db_conn, id_)
+        except NotFoundException:
+            return cls(id_)
+
     @classmethod
     def all(cls: type[BaseModelInstance],
             db_conn: DatabaseConnection) -> list[BaseModelInstance]:
@@ -465,7 +503,7 @@ class BaseModel(Generic[BaseModelId]):
                                       values)
         if not isinstance(self.id_, str):
             self.id_ = cursor.lastrowid  # type: ignore[assignment]
-        self._cache()
+        self.cache()
         for attr_name in self.to_save_versioned:
             getattr(self, attr_name).save(db_conn)
         for table, column, attr_name, key_index in self.to_save_relations: