diff --git a/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py b/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py index c4866fcecc..f7d614e73a 100644 --- a/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py +++ b/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py @@ -449,26 +449,31 @@ async def _fetch_zip_from_url_async( def _download_and_parse() -> dict[str, list[dict[str, Any]]]: zip_path: Path temp_to_clean: Path | None = None - if cache and cache_path.exists(): - zip_path = cache_path - else: - if cache: - cache_dir.mkdir(parents=True, exist_ok=True) + try: + if cache and cache_path.exists(): zip_path = cache_path else: - with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp: - zip_path = Path(tmp.name) + if cache: + cache_dir.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + delete=False, + dir=cache_dir, + suffix=".zip.part", + ) as tmp: + zip_path = Path(tmp.name) + else: + with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp: + zip_path = Path(tmp.name) temp_to_clean = zip_path - logger.info(f"Downloading zip archive from {source}") - with requests.get(source, stream=True) as response: - response.raise_for_status() - with zip_path.open("wb") as fh: - for chunk in response.iter_content(chunk_size=1 << 16): - if chunk: - fh.write(chunk) + logger.info(f"Downloading zip archive from {source}") + with requests.get(source, stream=True) as response: + response.raise_for_status() + with zip_path.open("wb") as fh: + for chunk in response.iter_content(chunk_size=1 << 16): + if chunk: + fh.write(chunk) - try: results: dict[str, list[dict[str, Any]]] = {} with zipfile.ZipFile(zip_path) as zf: members = set(zf.namelist()) @@ -485,6 +490,9 @@ def _download_and_parse() -> dict[str, list[dict[str, Any]]]: "list[dict[str, Any]]", self.FILE_TYPE_HANDLERS[file_type]["read"](text), ) + if cache and temp_to_clean is not None: + zip_path.replace(cache_path) + temp_to_clean = None return results finally: if temp_to_clean is not None: diff --git a/tests/unit/datasets/test_remote_dataset_loader.py b/tests/unit/datasets/test_remote_dataset_loader.py index 1978e48425..7d74fb7714 100644 --- a/tests/unit/datasets/test_remote_dataset_loader.py +++ b/tests/unit/datasets/test_remote_dataset_loader.py @@ -256,6 +256,43 @@ async def test_caches_zip_on_disk(self, tmp_path, monkeypatch): cached = list((tmp_path / "seed-prompt-entries").glob("*.zip")) assert len(cached) == 1 + async def test_interrupted_download_does_not_poison_cache(self, tmp_path, monkeypatch): + monkeypatch.setattr( + "pyrit.datasets.seed_datasets.remote.remote_dataset_loader.DB_DATA_PATH", + tmp_path, + ) + zip_bytes = self._make_zip_bytes({"x.json": '[{"k": "v"}]'}) + + def interrupted_chunks(): + yield zip_bytes[:20] + raise RuntimeError("connection dropped") + + interrupted_response = self._mock_streaming_response(b"") + interrupted_response.iter_content.return_value = interrupted_chunks() + successful_response = self._mock_streaming_response(zip_bytes) + + with patch( + "pyrit.datasets.seed_datasets.remote.remote_dataset_loader.requests.get", + side_effect=[interrupted_response, successful_response], + ) as mock_get: + loader = ConcreteRemoteLoader() + with pytest.raises(RuntimeError, match="connection dropped"): + await loader._fetch_zip_from_url_async(source=self.SOURCE, inner_files=["x.json"], cache=True) + + cache_dir = tmp_path / "seed-prompt-entries" + assert list(cache_dir.iterdir()) == [] + + result = await loader._fetch_zip_from_url_async( + source=self.SOURCE, + inner_files=["x.json"], + cache=True, + ) + + assert result == {"x.json": [{"k": "v"}]} + assert mock_get.call_count == 2 + assert len(list(cache_dir.glob("*.zip"))) == 1 + assert not list(cache_dir.glob("*.part")) + async def test_cache_false_does_not_persist_zip(self, tmp_path, monkeypatch): monkeypatch.setattr( "pyrit.datasets.seed_datasets.remote.remote_dataset_loader.DB_DATA_PATH",