Spaces:
Sleeping
Sleeping
| 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() | |