From a7ecfa8a410fa05ff7b91e6f6c751649a8f38237 Mon Sep 17 00:00:00 2001 From: CaptainOfHacks <39195263+CaptainOfHacks@users.noreply.github.com> Date: Sat, 30 Dec 2023 13:34:43 +0200 Subject: [PATCH] Upgrade state manager component - update documentation for state manager - enable unit tests for state manager - update logic of state manager component - update mongomock with GridFS support version --- .../database/adapters/gridfs_storage.py | 47 ++++++++++++++++--- .../services/object_state_manager.py | 15 ++++-- requirements.dev.txt | 2 +- tests/unit/backend/conftest.py | 1 - .../state_manager/_test_state_manager.py | 17 ------- .../state_manager/test_state_manager.py | 22 +++++++++ 6 files changed, 76 insertions(+), 28 deletions(-) delete mode 100644 tests/unit/backend/state_manager/_test_state_manager.py create mode 100644 tests/unit/backend/state_manager/test_state_manager.py diff --git a/mapping_workbench/backend/database/adapters/gridfs_storage.py b/mapping_workbench/backend/database/adapters/gridfs_storage.py index 6d8623b07..ac3a1d0e6 100644 --- a/mapping_workbench/backend/database/adapters/gridfs_storage.py +++ b/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) diff --git a/mapping_workbench/backend/state_manager/services/object_state_manager.py b/mapping_workbench/backend/state_manager/services/object_state_manager.py index c7d003b2e..c8e70ac73 100644 --- a/mapping_workbench/backend/state_manager/services/object_state_manager.py +++ b/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) diff --git a/requirements.dev.txt b/requirements.dev.txt index 9b4c8138c..a764e09fe 100644 --- a/requirements.dev.txt +++ b/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 \ No newline at end of file +mongomock-motor==0.0.26 \ No newline at end of file diff --git a/tests/unit/backend/conftest.py b/tests/unit/backend/conftest.py index 201f24e91..26cb22e1a 100644 --- a/tests/unit/backend/conftest.py +++ b/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 diff --git a/tests/unit/backend/state_manager/_test_state_manager.py b/tests/unit/backend/state_manager/_test_state_manager.py deleted file mode 100644 index 262cbd138..000000000 --- a/tests/unit/backend/state_manager/_test_state_manager.py +++ /dev/null @@ -1,17 +0,0 @@ -import pytest - -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(): - 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) - new_fake_object_state = await load_object_state(fake_object_state_id, FakeObjectState) - assert new_fake_object_state is None \ No newline at end of file diff --git a/tests/unit/backend/state_manager/test_state_manager.py b/tests/unit/backend/state_manager/test_state_manager.py new file mode 100644 index 000000000..9e81d050f --- /dev/null +++ b/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