from __future__ import annotations import os import sys import tempfile import unittest import zipfile from pathlib import Path PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__)) if PROJECT_ROOT not in sys.path: sys.path.insert(0, PROJECT_ROOT) from core.zip_handler import ( # noqa: E402 apply_ai_changes, build_zip_ai_context, extract_uploaded_zip, recompress_project_zip, ) class ZipHandlerTests(unittest.TestCase): def _make_sample_zip(self, tmpdir: str) -> str: zip_path = os.path.join(tmpdir, "sample_project.zip") with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: zf.writestr("sample_project/app.py", "print('hello')\n") zf.writestr("sample_project/core/utils.py", "def add(a, b):\n return a + b\n") zf.writestr("sample_project/README.md", "# Sample\n") return zip_path def test_extract_build_context_and_recompress(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: zip_path = self._make_sample_zip(tmpdir) bundle = extract_uploaded_zip(zip_path, workspace=tmpdir, project_id="ziptest") self.assertTrue(os.path.isdir(bundle["project_root"])) self.assertEqual(bundle["archive_root_prefix"], "sample_project") context = build_zip_ai_context(bundle["project_root"], focus_files=["app.py"]) self.assertIn("ZIP Intelligence", context) self.assertIn("### FILE: app.py", context) self.assertIn("core/utils.py", context) result = apply_ai_changes( bundle["project_root"], """```python # FILE: app.py print('zip intelligence') ``` ```python # FILE: core/utils.py def add(a, b): return a + b + 1 ``` """, ) self.assertTrue(result["applied"]) self.assertEqual(result["strategy"], "file_blocks") self.assertIn("app.py", result["modified_files"]) app_text = Path(bundle["project_root"]).joinpath("app.py").read_text(encoding="utf-8") self.assertIn("zip intelligence", app_text) out_zip = recompress_project_zip( bundle["project_root"], os.path.join(tmpdir, "out"), original_zip_name=os.path.basename(zip_path), archive_root_prefix=bundle["archive_root_prefix"], ) self.assertTrue(os.path.exists(out_zip)) with zipfile.ZipFile(out_zip, "r") as zf: names = set(zf.namelist()) self.assertIn("sample_project/app.py", names) self.assertIn("sample_project/core/utils.py", names) def test_apply_ai_changes_returns_error_for_empty_output(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: zip_path = self._make_sample_zip(tmpdir) bundle = extract_uploaded_zip(zip_path, workspace=tmpdir, project_id="ziptest_empty") result = apply_ai_changes(bundle["project_root"], "") self.assertFalse(result["applied"]) self.assertIn("empty generated output", " ".join(result["errors"])) if __name__ == "__main__": unittest.main()