diff --git a/housekeeper/store/api/handlers/update.py b/housekeeper/store/api/handlers/update.py index c7e5c93b..911554ca 100644 --- a/housekeeper/store/api/handlers/update.py +++ b/housekeeper/store/api/handlers/update.py @@ -1,19 +1,14 @@ """ This module handles updating entries in the store/database. """ -from datetime import datetime import logging +from datetime import datetime from typing import List -from sqlalchemy.orm import Session - - from housekeeper.store.api.handlers.base import BaseHandler -from housekeeper.store.filters.archive_filters import ( - apply_archive_filter, - ArchiveFilter, -) -from housekeeper.store.models import Archive +from housekeeper.store.filters.archive_filters import ArchiveFilter, apply_archive_filter +from housekeeper.store.models import Archive, File +from sqlalchemy.orm import Session LOG = logging.getLogger(__name__) @@ -59,3 +54,13 @@ def update_finished_retrieval_task(self, retrieval_task_id: int) -> None: ).all() for archive in completed_archives: self.update_retrieval_time_stamp(archive=archive) + + @staticmethod + def update_retrieval_task_id(archive: Archive, retrieval_task_id: int): + """Sets the retrieval_task_id in the Archive entry for the provided file.""" + archive.retrieval_task_id: int = retrieval_task_id + + @staticmethod + def update_archiving_task_id(archive: Archive, archiving_task_id: int): + """Sets the archiving_task_id in the Archive entry for the provided file.""" + archive.archiving_task_id: int = archiving_task_id diff --git a/tests/conftest.py b/tests/conftest.py index 4061f8cb..12607bfe 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -78,6 +78,12 @@ def fixture_archiving_task_id() -> int: return 1234 +@pytest.fixture(scope="function", name="new_archiving_task_id") +def fixture_new_archiving_task_id() -> int: + """Return a new id of an archiving task.""" + return 1235 + + @pytest.fixture(scope="function", name="retrieval_task_id") def fixture_retrieval_task_id() -> int: """Return an id of a retrieval task.""" diff --git a/tests/store/handlers/test_updatehandler.py b/tests/store/handlers/test_updatehandler.py index c0bd2682..e6d74217 100644 --- a/tests/store/handlers/test_updatehandler.py +++ b/tests/store/handlers/test_updatehandler.py @@ -1,4 +1,5 @@ from datetime import datetime + from housekeeper import store from housekeeper.store import Store from housekeeper.store.models import Archive @@ -94,3 +95,32 @@ def test_update_finished_retrieval_task( assert archive.retrieved_at assert second_archive.retrieved_at assert archive.retrieved_at.minute == second_archive.retrieved_at.minute + + +def test_update_retrieval_task_id(archive: Archive, retrieval_task_id: int, populated_store: store): + """Tests updating the retrieved_at timestamp on a given archive when there already is one.""" + # GIVEN an Archive with no retrieval_task_id + assert not archive.retrieval_task_id + + # WHEN updating the retrieval task id + populated_store.update_retrieval_task_id(archive=archive, retrieval_task_id=retrieval_task_id) + + # THEN the retrieval task id should be set + assert archive.retrieval_task_id == retrieval_task_id + + +def test_update_archiving_task_id( + archive: Archive, new_archiving_task_id: int, populated_store: store +): + """Tests updating the archiving_task_od on a given archive when there already is one. + Necessary for retrieval and rearchiving.""" + # GIVEN an Archive with an old archiving_task_id + assert archive.archiving_task_id != new_archiving_task_id + + # WHEN updating the archiving_task_id + populated_store.update_archiving_task_id( + archive=archive, archiving_task_id=new_archiving_task_id + ) + + # THEN the retrieval task id should be set + assert archive.archiving_task_id == new_archiving_task_id