fix: use partitions for stage-packages tracking

This localized fix addresses the issue of stage-package tracking failing in
the presence of partitions by looking for the primed files in each partition.

This solves the most common cases but doesn't fix the issue where a given file
might be present in multiple partitions, as this needs a deeper refactoring of
the state tracking (see discussion in #804).

Fixes #804

craft_parts/executor/part_handler.py

@@ -420,8 +420,9 @@ def _run_prime(
             and self._track_stage_packages
             and is_deb_based()
         ):
+            prime_dirs = list(self._part.prime_dirs.values())
             primed_stage_packages = _get_primed_stage_packages(
-                contents.files, prime_dir=self._part.prime_dir
+                contents.files, prime_dirs=prime_dirs
             )
         else:
             primed_stage_packages = set()
@@ -1132,11 +1133,16 @@ def _parts_with_overlay_in_step(step: Step, *, part_list: list[Part]) -> list[Pa
     return [p for p in oparts if states.get_step_state_path(p, step).exists()]
 
 
-def _get_primed_stage_packages(snap_files: set[str], *, prime_dir: Path) -> set[str]:
+def _get_primed_stage_packages(
+    snap_files: set[str], *, prime_dirs: list[Path]
+) -> set[str]:
     primed_stage_packages: set[str] = set()
     for _snap_file in snap_files:
-        snap_file = str(prime_dir / _snap_file)
-        stage_package = read_origin_stage_package(snap_file)
-        if stage_package:
-            primed_stage_packages.add(stage_package)
+        for prime_dir in prime_dirs:
+            snap_file = prime_dir / _snap_file
+            if not snap_file.exists():
+                continue
+            stage_package = read_origin_stage_package(str(snap_file))
+            if stage_package:
+                primed_stage_packages.add(stage_package)
     return primed_stage_packages

tests/integration/features/partitions/lifecycle/test_lifecycle.py

@@ -302,3 +302,39 @@ def test_clean_all_parts(self, step, foo_files, bar_files):
 
 class TestUpdating(test_lifecycle.TestUpdating):
     """Run all updating tests with partitions enabled."""
+
+
+def test_track_stage_packages_with_partitions(new_dir):
+    partitions = ["default", "binaries", "docs"]
+    parts_yaml = textwrap.dedent(
+        """
+            parts:
+              foo:
+                plugin: nil
+                stage-packages:
+                  - hello
+                organize:
+                  usr/bin: (binaries)/
+                  usr/share/doc: (docs)/
+            """
+    )
+    parts = yaml.safe_load(parts_yaml)
+
+    lifecycle = craft_parts.LifecycleManager(
+        parts,
+        application_name="test_track_stage_packages_with_partitions",
+        cache_dir=new_dir,
+        partitions=partitions,
+        track_stage_packages=True,
+    )
+
+    actions = lifecycle.plan(Step.PRIME)
+
+    with lifecycle.action_executor() as ctx:
+        ctx.execute(actions)
+
+    packages = lifecycle.get_primed_stage_packages(part_name="foo")
+    assert packages is not None
+
+    name_only = [p.split("=")[0] for p in packages]
+    assert "hello" in name_only

tests/unit/executor/test_part_handler.py

@@ -200,6 +200,7 @@ def test_run_prime(self, new_dir, mocker):
             return_value=StepContents({"file"}, {"dir"}),
         )
         mocker.patch("os.getxattr", return_value=b"pkg")
+        mocker.patch("pathlib.Path.exists", return_value=True)
 
         ovmgr = OverlayManager(
             project_info=self._project_info, part_list=[self._part], base_layer_dir=None