Spaces:
Sleeping
Sleeping
File size: 3,213 Bytes
df6cd5e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | 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()
|