diff --git a/strix/interface/utils.py b/strix/interface/utils.py index 26af96ef3..8ae51a870 100644 --- a/strix/interface/utils.py +++ b/strix/interface/utils.py @@ -1501,6 +1501,10 @@ def rewrite_localhost_targets(targets_info: list[dict[str, Any]], host_gateway: #: from wherever they happen to live on the host. API_SPEC_WORKSPACE_SUBDIR = "api-specs" +# Keep ownership of downloads until staging succeeds. TemporaryDirectory also +# removes pending downloads at normal process exit if setup is abandoned. +_fetched_collections: dict[str, tempfile.TemporaryDirectory[str]] = {} + def write_fetched_collection(collection: dict[str, Any], collection_uid: str) -> str: """Write a collection fetched from the Postman API to a local file. @@ -1512,9 +1516,14 @@ def write_fetched_collection(collection: dict[str, Any], collection_uid: str) -> staging.mkdir(parents=True, exist_ok=True) # One collection can be targeted with multiple environments in the same run. # Keep each resolved snapshot independent, including across concurrent runs. - snapshot = Path(tempfile.mkdtemp(prefix="collection-", dir=staging)) - path = snapshot / f"{sanitize_name(collection_uid)}.postman_collection.json" - write_secret_text(path, json.dumps(collection, indent=2)) + snapshot = tempfile.TemporaryDirectory(prefix="collection-", dir=staging) + path = Path(snapshot.name) / f"{sanitize_name(collection_uid)}.postman_collection.json" + try: + write_secret_text(path, json.dumps(collection, indent=2)) + except BaseException: + snapshot.cleanup() + raise + _fetched_collections[str(path)] = snapshot return str(path) @@ -1544,9 +1553,22 @@ def stage_api_specs(targets_info: list[dict[str, Any]], run_name: str) -> list[d count += 1 name = f"{stem}-{count}{suffix}" used.add(name) - shutil.copy2(source, staging / name) + destination = staging / name + if source.resolve() != destination.resolve(): + shutil.copy2(source, destination) details["workspace_path"] = f"/workspace/{API_SPEC_WORKSPACE_SUBDIR}/{name}" + # Release only downloads this process created, after every copy succeeds. + # Retain the per-run staged files: resume still needs those snapshots. + for target in specs: + details = target["details"] + downloaded = str(details["target_spec"]) + snapshot = _fetched_collections.get(downloaded) + if snapshot is not None: + details["target_spec"] = str(staging / Path(details["workspace_path"]).name) + snapshot.cleanup() + del _fetched_collections[downloaded] + return [ { "source_path": str(staging), diff --git a/tests/test_postman_environment_handoff.py b/tests/test_postman_environment_handoff.py index 8ae89140b..52bba6108 100644 --- a/tests/test_postman_environment_handoff.py +++ b/tests/test_postman_environment_handoff.py @@ -5,6 +5,7 @@ from __future__ import annotations import argparse import copy import json +import shutil from pathlib import Path from types import SimpleNamespace from typing import Any @@ -162,3 +163,37 @@ def test_references_in_collection_values_find_environment_dependencies( variables = postman_variables(json.loads(Path(target["details"]["target_spec"]).read_text())) assert variables["token"] == "{{accessToken}}" # noqa: S105 - unresolved template assert variables["accessToken"] == "scan-only-token" + + +@pytest.mark.usefixtures("postman_api") +def test_staging_releases_the_download_snapshot_and_keeps_the_run_copy() -> None: + targets = _targets("postman://collection-1?env=env-1") + downloaded = Path(targets[0]["details"]["target_spec"]) + (source,) = stage_api_specs(targets, "snapshot-cleanup") + + assert not downloaded.parent.exists() + staged = Path(targets[0]["details"]["target_spec"]) + assert staged.parent == Path(source["source_path"]) + assert postman_variables(json.loads(staged.read_text()))["tenant"] == "tenant-one" + # Repeating setup uses the retained run snapshot rather than a removed input. + assert stage_api_specs(targets, "snapshot-cleanup") == [source] + + +@pytest.mark.usefixtures("postman_api") +def test_failed_staging_keeps_downloads_available_for_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + targets = _targets("postman://collection-1?env=env-1") + downloaded = Path(targets[0]["details"]["target_spec"]) + copy_file = shutil.copy2 + + def fail_copy(*_args: Any, **_kwargs: Any) -> None: + raise OSError("staging unavailable") + + monkeypatch.setattr(shutil, "copy2", fail_copy) + with pytest.raises(OSError, match="staging unavailable"): + stage_api_specs(targets, "failed-staging") + assert downloaded.is_file() + monkeypatch.setattr(shutil, "copy2", copy_file) + stage_api_specs(targets, "failed-staging") + assert not downloaded.parent.exists()