Functionality for updating archiving/retrieval task ids

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 = 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 = archiving_task_id

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."""

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