From f5b50b6263b22fadb18eeb5f0b883c7ec6cbb524 Mon Sep 17 00:00:00 2001
From: cecinestpasunepipe
 <110607403+cecinestpasunepipe@users.noreply.github.com>
Date: Mon, 1 Jul 2024 09:58:21 +0200
Subject: [PATCH] Fix shard issue (#175)

(DIS-3251)
---
 acquire/acquire.py |  4 ++--
 tests/test_gui.py  | 48 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_gui.py

diff --git a/acquire/acquire.py b/acquire/acquire.py
index ee7ee8bc..2fc395fd 100644
--- a/acquire/acquire.py
+++ b/acquire/acquire.py
@@ -2167,7 +2167,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))
 
@@ -2180,7 +2180,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:
diff --git a/tests/test_gui.py b/tests/test_gui.py
new file mode 100644
index 00000000..bd225d55
--- /dev/null
+++ b/tests/test_gui.py
@@ -0,0 +1,48 @@
+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(
+    "num_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, num_children: int, skip_parent: bool, auto_upload: bool, expected_shards: list[int]
+):
+    def list_children() -> list[Target]:
+        yield from [mock_target] * num_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=True, skip_parent=skip_parent, start_time=0)
+    acquire_children_and_targets(mock_target, args)
+    assert shards == expected_shards