FIX Avoid poisoned dataset ZIP caches

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:

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",