From 5b43a4b39655304cf18249211e66b139f1f35e5f Mon Sep 17 00:00:00 2001
From: Sebin Sunny <sebin.sunny@aiven.io>
Date: Fri, 1 Mar 2024 11:01:17 +1100
Subject: [PATCH] feat: Add progress_callback to snapshot method

Updated the snapshot method to include a progress_callback for better tracking of operation. This allows for real-time progress tracking during snapshot operations, improving logging and debuggability for delta backup
---
 rohmu/delta/snapshot.py     | 16 ++++++++++++-
 test/delta/test_snapshot.py | 48 +++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/rohmu/delta/snapshot.py b/rohmu/delta/snapshot.py
index 7f72004d..395bc145 100644
--- a/rohmu/delta/snapshot.py
+++ b/rohmu/delta/snapshot.py
@@ -218,7 +218,13 @@ def _snapshot_add_missing_files(self, *, src_files: Sequence[Path], dst_files: S
             changes += 1
         return changes
 
-    def snapshot(self, *, progress: Optional[Progress] = None, reuse_old_snapshotfiles: bool = True) -> int:
+    def snapshot(
+        self,
+        *,
+        progress: Optional[Progress] = None,
+        reuse_old_snapshotfiles: bool = True,
+        progress_callback: Optional[Callable[[str, bool], None]] = None,
+    ) -> int:
         assert self.lock.locked()
 
         if progress is None:
@@ -258,14 +264,20 @@ def snapshot(self, *, progress: Optional[Progress] = None, reuse_old_snapshotfil
         # Create missing directories
         changes = self._snapshot_create_missing_directories(src_dirs=src_dirs, dst_dirs=dst_dirs)
         progress.add_success()
+        if progress_callback:
+            progress_callback("Creating missing directories", True)
 
         # Remove extra files
         changes += self._snapshot_remove_extra_files(src_files=src_files, dst_files=dst_files)
         progress.add_success()
+        if progress_callback:
+            progress_callback("Removing extra files", True)
 
         # Add missing files
         changes += self._snapshot_add_missing_files(src_files=src_files, dst_files=dst_files)
         progress.add_success()
+        if progress_callback:
+            progress_callback("Adding missing files", True)
 
         # We COULD also remove extra directories, but it is not
         # probably really worth it and due to ignored files it
@@ -294,6 +306,8 @@ def _result_cb(*, map_in: Any, map_out: SnapshotFile) -> bool:
             self._add_snapshotfile(map_out)
             assert progress is not None
             progress.add_success()
+            if progress_callback:
+                progress_callback("Processing and hashing snapshot files", True)
             return True
 
         changes += len(snapshotfiles)
diff --git a/test/delta/test_snapshot.py b/test/delta/test_snapshot.py
index d10df2b6..bbe9df20 100644
--- a/test/delta/test_snapshot.py
+++ b/test/delta/test_snapshot.py
@@ -202,3 +202,51 @@ def fake_open_for_reading(self: SnapshotFile, path: Path) -> SizeLimitedFile:
         with patch("rohmu.delta.snapshot.Path.stat", side_effect=FileNotFoundError):
             with pytest.raises(FileNotFoundError):
                 snapshotter.snapshot(progress=Progress(), reuse_old_snapshotfiles=True)
+
+
+@pytest.mark.timeout(2)
+def test_snapshot_without_create_samples(snapshotter_creator: Callable[..., SnapshotterWithDefaults]) -> None:
+    snapshotter = snapshotter_creator()
+    callback_messages: list[str] = []
+
+    def progress_callback(message: str, success: bool) -> None:
+        callback_message = f"{message}: {'Success' if success else 'Failure'}"
+        callback_messages.append(callback_message)
+
+    samples: dict[Union[str, Path], str] = {
+        "foo": "foobar",
+        "foo2": "foobar",
+        "foobig": "foobar" * EMBEDDED_FILE_SIZE,
+        "foobig2": "foobar" * EMBEDDED_FILE_SIZE,
+    }
+
+    src = snapshotter.src
+    for file_name, body in samples.items():
+        (src / file_name).write_text(body)
+
+    with snapshotter.lock:
+        progress = Progress()
+        changes_detected = snapshotter.snapshot(progress=progress, progress_callback=progress_callback)
+        assert changes_detected > 0
+
+        ss1 = snapshotter.get_snapshot_state()
+
+        progress = Progress()
+        assert snapshotter.snapshot(progress=progress, progress_callback=progress_callback) == 0
+
+        ss2 = snapshotter.get_snapshot_state()
+        assert ss1 == ss2
+        expected_messages = [
+            "Creating missing directories: Success",
+            "Removing extra files: Success",
+            "Adding missing files: Success",
+            "Processing and hashing snapshot files: Success",
+            "Processing and hashing snapshot files: Success",
+            "Processing and hashing snapshot files: Success",
+            "Processing and hashing snapshot files: Success",
+            "Creating missing directories: Success",
+            "Removing extra files: Success",
+            "Adding missing files: Success",
+        ]
+
+        assert callback_messages == expected_messages