File size: 5,308 Bytes
974567e | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | 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()
|