Fix shard issue (DIS-3251)

acquire/acquire.py

@@ -2158,7 +2158,7 @@ def acquire_children_and_targets(target: Target, args: argparse.Namespace) -> No
     if (args.children and not args.skip_parent) or not args.children:
         total_targets += 1
         counter += 1
-        acquire_gui.shard = (progress_limit // total_targets) * counter
+        acquire_gui.shard = int((progress_limit / total_targets) * counter)
         try:
             files.extend(acquire_target(target, args, args.start_time))
 
@@ -2171,7 +2171,7 @@ def acquire_children_and_targets(target: Target, args: argparse.Namespace) -> No
     if args.children:
         for child in target.list_children():
             counter += 1
-            acquire_gui.shard = (100 // total_targets) * counter
+            acquire_gui.shard = int((progress_limit / total_targets) * counter)
             try:
                 child_target = load_child(target, child.path)
             except Exception:

tests/test_gui.py

@@ -0,0 +1,49 @@
+from argparse import Namespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+from dissect.target import Target
+
+from acquire.acquire import acquire_children_and_targets
+from acquire.gui import GUI
+
+
+@pytest.mark.parametrize(
+    "children, skip_parent, auto_upload, expected_shards",
+    [
+        (0, False, False, [90]),  # 90 (default, leaves 10% for final 'step')
+        (0, False, True, [50]),  # 50% till upload (upload progresses in plugin)
+        (1, False, False, [45, 90]),  # two children to 90%
+        (1, True, False, [90]),  # without parent, it's just one target - so 90%
+        (1, False, True, [25, 50]),  # two children to 50%
+        (1, True, True, [50]),  # one till upload (50%)
+        (2, False, False, [30, 60, 90]),  # two children + parent till 90%
+        (2, False, True, [16, 33, 50]),  # two children + parent till 50%
+        (50, False, True, list(range(0, 51))),  # Should not be zero filled...
+    ],
+)
+@patch("acquire.gui.base.Stub", spec=True)
+@patch("acquire.acquire.acquire_target", create=True)
+def test_gui(
+    mock_target: Target, gui: GUI, children: int, skip_parent: bool, auto_upload: bool, expected_shards: list[int]
+):
+    def list_children():
+        return children
+
+    mock_target.list_children = list_children
+
+    class Diagnostic_GUI(MagicMock):
+        @property
+        def shard(self) -> int:
+            return 0
+
+        @shard.setter
+        def shard(self, shard: int) -> None:
+            shards.append(shard)
+
+    GUI.__new__ = lambda x: Diagnostic_GUI()
+    shards = []
+    args = Namespace(child=False, auto_upload=auto_upload, children=children, skip_parent=skip_parent, start_time=0)
+    children = [mock_target] * children
+    acquire_children_and_targets(mock_target, args)
+    assert shards == expected_shards