Upgrade state manager component

mapping_workbench/backend/database/adapters/gridfs_storage.py

@@ -1,41 +1,76 @@
 from io import BytesIO
 import gzip
+from typing import Optional
 
 from motor.motor_asyncio import AsyncIOMotorGridFSBucket, AsyncIOMotorDatabase
 
 
 class AsyncGridFSStorage:
+    """
+    This class is a wrapper for the AsyncIOMotorGridFSBucket class.
+    """
     _mongo_database: AsyncIOMotorDatabase = None
 
     @classmethod
     def set_mongo_database(cls, mongo_database: AsyncIOMotorDatabase):
+        """
+        Sets the mongo database to use for the gridfs storage.
+        :param mongo_database: The mongo database to use for the gridfs storage.
+        :return: None
+        """
         cls._mongo_database = mongo_database
 
     @classmethod
     def get_mongo_database(cls) -> AsyncIOMotorDatabase:
+        """
+        Gets the mongo database to use for the gridfs storage.
+        :return: The mongo database to use for the gridfs storage.
+        """
         if cls._mongo_database is None:
             from mapping_workbench.backend.database.adapters.mongodb import DB
             cls._mongo_database = DB.get_database()
         return cls._mongo_database
 
     @classmethod
-    async def upload_file(cls, file_id: str, file_content: str):
+    async def upload_file(cls, file_name: str, file_content: str) -> str:
+        """
+        Uploads a file to the gridfs storage.
+        :param file_name: The name of the file to upload.
+        :param file_content: The content of the file to upload.
+        :return: The id of the uploaded file.
+        """
         mongo_db = cls.get_mongo_database()
         grid_fs = AsyncIOMotorGridFSBucket(mongo_db)
         compressed_data = gzip.compress(file_content.encode("utf-8"))
-        await grid_fs.upload_from_stream(file_id, compressed_data)
+        file_id = await grid_fs.upload_from_stream(file_name, compressed_data)
+        return file_id
 
     @classmethod
-    async def download_file(cls, file_id: str) -> str:
+    async def download_file(cls, file_id: str) -> Optional[str]:
+        """
+        Downloads a file from the gridfs storage.
+        :param file_id: The id of the file to download.
+        :return: The content of the downloaded file.
+        """
         mongo_db = cls.get_mongo_database()
         grid_fs = AsyncIOMotorGridFSBucket(mongo_db)
         tmp_stream = BytesIO()
-        await grid_fs.download_to_stream_by_name(file_id, tmp_stream)
-        compressed_data = tmp_stream.read()
-        return gzip.decompress(compressed_data).decode("utf-8")
+        try:
+            await grid_fs.download_to_stream(file_id, tmp_stream)
+            compressed_data = tmp_stream.getvalue()
+            result_data = gzip.decompress(compressed_data).decode("utf-8")
+        except Exception:
+            result_data = None
+        tmp_stream.close()
+        return result_data
 
     @classmethod
     async def delete_file(cls, file_id: str):
+        """
+        Deletes a file from the gridfs storage.
+        :param file_id: The id of the file to delete.
+        :return: None
+        """
         mongo_db = cls.get_mongo_database()
         grid_fs = AsyncIOMotorGridFSBucket(mongo_db)
         await grid_fs.delete(file_id)

mapping_workbench/backend/state_manager/services/object_state_manager.py

@@ -9,23 +9,32 @@
 async def save_object_state(object_state: ObjectState) -> str:
     """
     Saves the state of an object to the database and return the saved state id.
+    :param object_state: The state of the object to save.
+    :return: The id of the saved state.
     """
     state_content_dump = object_state.model_dump_json()
-    state_id = sha1(state_content_dump.encode("utf-8")).hexdigest()
-    await AsyncGridFSStorage.upload_file(state_id, state_content_dump)
-    return state_id
+    file_name = str(sha1(state_content_dump.encode("utf-8")).hexdigest())
+    grids_fs_state_id = await AsyncGridFSStorage.upload_file(file_name, state_content_dump)
+    return grids_fs_state_id
 
 
 async def load_object_state(state_id: str, object_class: Type[ObjectStateType]) -> Optional[ObjectStateType]:
     """
     Loads the state of an object from the database.
+    :param state_id: The id of the state to load.
+    :param object_class: The class of the object to load.
+    :return: The loaded object.
     """
     state_content_dump = await AsyncGridFSStorage.download_file(state_id)
+    if state_content_dump is None:
+        return None
     return object_class(**json.loads(state_content_dump))
 
 
 async def delete_object_state(state_id: str):
     """
     Deletes the state of an object from the database.
+    :param state_id: The id of the state to delete.
+    :return: None
     """
     return await AsyncGridFSStorage.delete_file(state_id)

requirements.dev.txt

@@ -8,4 +8,4 @@ tox~=3.24.5
 tox-pytest-summary~=0.1.2
 mongomock==4.1.2
 pytest-asyncio~=0.21.1
-mongomock-motor==0.0.25
+mongomock-motor==0.0.26

tests/unit/backend/conftest.py

@@ -2,7 +2,6 @@
 
 from mongomock_motor import AsyncMongoMockClient
 from mapping_workbench.backend.core.services.project_initilisers import init_project_models
-from mapping_workbench.backend.database.adapters.gridfs_storage import AsyncGridFSStorage
 
 from mapping_workbench.backend.database.adapters.gridfs_storage import AsyncGridFSStorage
 

tests/unit/backend/state_manager/test_state_manager.py

@@ -0,0 +1,22 @@
+import pytest
+from gridfs import NoFile
+from mongomock_motor import enabled_gridfs_integration
+
+from mapping_workbench.backend.state_manager.services.object_state_manager import save_object_state, load_object_state, \
+    delete_object_state
+from tests.fakes.fake_state_object import FakeObjectState
+
+
+@pytest.mark.asyncio
+async def test_object_state_manager():
+    with enabled_gridfs_integration():
+        fake_object_state = FakeObjectState(name="Test1", object_data="Test2")
+        fake_object_state_id = await save_object_state(fake_object_state)
+        new_fake_object_state = await load_object_state(fake_object_state_id, FakeObjectState)
+        assert new_fake_object_state.name == fake_object_state.name
+        assert new_fake_object_state.object_data == fake_object_state.object_data
+        await delete_object_state(fake_object_state_id)
+        with pytest.raises(NoFile):
+            await delete_object_state(fake_object_state_id)
+        new_fake_object_state = await load_object_state(fake_object_state_id, FakeObjectState)
+        assert new_fake_object_state is None