Initial project import

2026-09-18 03:46:58 +08:00
commit b600da4b76
209 changed files with 45904 additions and 0 deletions
@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+from PIL import Image
+
+from services.image_storage_service import ImageStorageService
+
+
+def png_bytes() -> bytes:
+    path = Path(tempfile.gettempdir()) / "chatgpt2api-test-image.png"
+    Image.new("RGB", (2, 2), color=(255, 0, 0)).save(path, format="PNG")
+    return path.read_bytes()
+
+
+class FakeWebDAVClient:
+    uploaded: dict[str, bytes] = {}
+    deleted: list[str] = []
+
+    def __init__(self, _settings):
+        pass
+
+    def put(self, rel: str, payload: bytes) -> str:
+        self.uploaded[rel] = payload
+        return f"https://dav.example.test/{rel}"
+
+    def get(self, rel: str) -> bytes:
+        return self.uploaded[rel]
+
+    def delete(self, rel: str) -> bool:
+        self.deleted.append(rel)
+        self.uploaded.pop(rel, None)
+        return True
+
+    def test(self) -> dict[str, object]:
+        self.put(".chatgpt2api_webdav_test.txt", b"chatgpt2api webdav test\n")
+        self.delete(".chatgpt2api_webdav_test.txt")
+        return {"ok": True, "status": 200, "error": None}
+
+
+class ImageStorageServiceTests(unittest.TestCase):
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.tmp.cleanup)
+        self.data_dir = Path(self.tmp.name)
+        self.images_dir = self.data_dir / "images"
+        self.settings = {
+            "enabled": False,
+            "mode": "local",
+            "webdav_url": "",
+            "webdav_username": "",
+            "webdav_password": "",
+            "webdav_root_path": "chatgpt2api/images",
+            "public_base_url": "",
+        }
+        self.config_patcher = mock.patch("services.image_storage_service.config")
+        self.mock_config = self.config_patcher.start()
+        self.addCleanup(self.config_patcher.stop)
+        self.mock_config.images_dir = self.images_dir
+        self.mock_config.base_url = "http://app.test"
+        self.mock_config.cleanup_old_images.return_value = 0
+        self.mock_config.get_image_storage_settings.side_effect = lambda: dict(self.settings)
+        FakeWebDAVClient.uploaded = {}
+        FakeWebDAVClient.deleted = []
+
+    def service(self) -> ImageStorageService:
+        return ImageStorageService(self.data_dir / "image_index.json")
+
+    def test_local_mode_saves_to_local_directory(self):
+        stored = self.service().save(png_bytes(), "http://app.test")
+
+        self.assertEqual(stored.storage, "local")
+        self.assertTrue((self.images_dir / stored.rel).is_file())
+        self.assertEqual(stored.url, f"http://app.test/images/{stored.rel}")
+
+    def test_webdav_mode_uploads_without_local_file(self):
+        self.settings.update({
+            "enabled": True,
+            "mode": "webdav",
+            "webdav_url": "https://dav.example.test",
+            "webdav_password": "secret",
+        })
+        with mock.patch("services.image_storage_service.WebDAVClient", FakeWebDAVClient):
+            stored = self.service().save(png_bytes(), "http://app.test")
+            payload = self.service().get_bytes(stored.rel)
+
+        self.assertEqual(stored.storage, "webdav")
+        self.assertFalse((self.images_dir / stored.rel).exists())
+        self.assertIn(stored.rel, FakeWebDAVClient.uploaded)
+        self.assertEqual(payload, FakeWebDAVClient.uploaded[stored.rel])
+
+    def test_list_items_ignores_non_image_files(self):
+        image = png_bytes()
+        image_path = self.images_dir / "2026" / "05" / "07" / "sample.png"
+        image_path.parent.mkdir(parents=True, exist_ok=True)
+        image_path.write_bytes(image)
+        (self.images_dir / ".DS_Store").write_text("not an image", encoding="utf-8")
+        (self.images_dir / "2026" / ".DS_Store").write_text("not an image", encoding="utf-8")
+
+        items = self.service().list_items("http://app.test")
+
+        self.assertEqual([item["rel"] for item in items], ["2026/05/07/sample.png"])
+        self.assertEqual(items[0]["storage"], "local")
+
+    def test_both_mode_saves_to_local_and_webdav(self):
+        self.settings.update({
+            "enabled": True,
+            "mode": "both",
+            "webdav_url": "https://dav.example.test",
+            "webdav_password": "secret",
+            "public_base_url": "https://cdn.example.test/images",
+        })
+        with mock.patch("services.image_storage_service.WebDAVClient", FakeWebDAVClient):
+            stored = self.service().save(png_bytes(), "http://app.test")
+
+        self.assertEqual(stored.storage, "both")
+        self.assertTrue((self.images_dir / stored.rel).is_file())
+        self.assertIn(stored.rel, FakeWebDAVClient.uploaded)
+        self.assertEqual(stored.url, f"https://cdn.example.test/images/{stored.rel}")
+
+    def test_test_webdav_writes_and_deletes_probe_file(self):
+        self.settings.update({
+            "enabled": True,
+            "mode": "webdav",
+            "webdav_url": "https://dav.example.test",
+            "webdav_password": "secret",
+        })
+        with mock.patch("services.image_storage_service.WebDAVClient", FakeWebDAVClient):
+            result = self.service().test_webdav()
+
+        self.assertTrue(result["ok"])
+        self.assertIn(".chatgpt2api_webdav_test.txt", FakeWebDAVClient.deleted)
+
+
+if __name__ == "__main__":
+    unittest.main()