diff --git a/tools/modly-cli/agent.py b/tools/modly-cli/agent.py index e424aa25..c7f37248 100644 --- a/tools/modly-cli/agent.py +++ b/tools/modly-cli/agent.py @@ -86,16 +86,25 @@ def _request_json( def _download(url: str, dest: Path, *, timeout: float) -> int: - dest.parent.mkdir(parents=True, exist_ok=True) + temporary_path: Path | None = None try: - with urllib.request.urlopen(url, timeout=timeout) as resp, dest.open("wb") as fh: + dest.parent.mkdir(parents=True, exist_ok=True) + # Keep the previous export intact until the download has finished. + # A sibling temporary file allows replacement on the same filesystem. + with urllib.request.urlopen(url, timeout=timeout) as resp, tempfile.NamedTemporaryFile( + dir=dest.parent, prefix=".modly-download-", suffix=".tmp", delete=False + ) as fh: + temporary_path = Path(fh.name) total = 0 while True: chunk = resp.read(1024 * 1024) if not chunk: - return total + break fh.write(chunk) total += len(chunk) + # Close the temporary file before replacing it, including on Windows. + os.replace(temporary_path, dest) + return total except urllib.error.HTTPError as exc: detail = exc.read().decode("utf-8", errors="replace") raise ModlyCliError(f"HTTP {exc.code} while downloading {url}: {detail}", code=f"HTTP_{exc.code}", http_status=exc.code) from exc @@ -103,6 +112,9 @@ def _download(url: str, dest: Path, *, timeout: float) -> int: raise ModlyCliError(f"Cannot download {url}: {exc.reason}", code="DOWNLOAD_FAILED") from exc except OSError as exc: raise ModlyCliError(f"Cannot write to {dest}: {exc}", code="WRITE_FAILED") from exc + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) def _multipart_form(fields: dict[str, str], file_field: str, file_path: Path) -> tuple[bytes, str]: diff --git a/tools/modly-cli/test_agent.py b/tools/modly-cli/test_agent.py index 9537ec1a..5598d3a0 100644 --- a/tools/modly-cli/test_agent.py +++ b/tools/modly-cli/test_agent.py @@ -11,7 +11,7 @@ from contextlib import redirect_stdout from pathlib import Path from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch MODULE_PATH = Path(__file__).with_name("agent.py") SPEC = importlib.util.spec_from_file_location("modly_agent", MODULE_PATH) @@ -28,6 +28,67 @@ def test_compact_json_is_one_line(self) -> None: self.assertEqual(buf.getvalue(), '{"nested":{"x":1},"ok":true}\n') +class DownloadTests(unittest.TestCase): + def test_success_replaces_destination_only_after_download(self) -> None: + with tempfile.TemporaryDirectory() as td: + dest = Path(td) / "mesh.glb" + dest.write_bytes(b"original") + chunks = iter([b"new ", b"mesh", b""]) + + def read(_size: int) -> bytes: + self.assertEqual(dest.read_bytes(), b"original") + return next(chunks) + + response = MagicMock() + response.__enter__.return_value.read.side_effect = read + with patch.object(agent.urllib.request, "urlopen", return_value=response): + self.assertEqual(agent._download("http://example.test/mesh", dest, timeout=1), 8) + self.assertEqual(dest.read_bytes(), b"new mesh") + self.assertEqual(list(Path(td).iterdir()), [dest]) + + def test_interrupted_download_preserves_destination_and_cleans_temporary_file(self) -> None: + for existing in (False, True): + for failure in (ConnectionResetError("connection lost"), KeyboardInterrupt()): + with self.subTest(existing=existing, failure=type(failure).__name__): + with tempfile.TemporaryDirectory() as td: + dest = Path(td) / "mesh.glb" + if existing: + dest.write_bytes(b"original") + response = MagicMock() + response.__enter__.return_value.read.side_effect = [b"partial", failure] + expected = KeyboardInterrupt if isinstance(failure, KeyboardInterrupt) else agent.ModlyCliError + with patch.object(agent.urllib.request, "urlopen", return_value=response): + with self.assertRaises(expected): + agent._download("http://example.test/mesh", dest, timeout=1) + if existing: + self.assertEqual(dest.read_bytes(), b"original") + else: + self.assertFalse(dest.exists()) + self.assertEqual(list(Path(td).iterdir()), [dest] if existing else []) + + def test_replace_failure_preserves_destination_and_cleans_temporary_file(self) -> None: + with tempfile.TemporaryDirectory() as td: + dest = Path(td) / "mesh.glb" + dest.write_bytes(b"original") + with ( + patch.object(agent.urllib.request, "urlopen", return_value=io.BytesIO(b"new mesh")), + patch.object(agent.os, "replace", side_effect=PermissionError("destination locked")), + ): + with self.assertRaises(agent.ModlyCliError) as ctx: + agent._download("http://example.test/mesh", dest, timeout=1) + self.assertEqual(ctx.exception.code, "WRITE_FAILED") + self.assertEqual(dest.read_bytes(), b"original") + self.assertEqual(list(Path(td).iterdir()), [dest]) + + def test_success_creates_destination_in_new_directory(self) -> None: + with tempfile.TemporaryDirectory() as td: + dest = Path(td) / "exports" / "mesh.glb" + with patch.object(agent.urllib.request, "urlopen", return_value=io.BytesIO(b"new mesh")): + self.assertEqual(agent._download("http://example.test/mesh", dest, timeout=1), 8) + self.assertEqual(dest.read_bytes(), b"new mesh") + self.assertEqual(list(dest.parent.iterdir()), [dest]) + + class CommandTests(unittest.TestCase): def test_status_combines_health_and_model(self) -> None: calls: list[tuple[str, str]] = []