diff --git a/acquire/acquire.py b/acquire/acquire.py index 5d172abd..3aa5b010 100644 --- a/acquire/acquire.py +++ b/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: diff --git a/tests/test_gui.py b/tests/test_gui.py new file mode 100644 index 00000000..55cd43d6 --- /dev/null +++ b/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