from __future__ import annotations import asyncio import hashlib import io import json import tarfile import time from pathlib import Path import httpx from imagefind.config import Settings from imagefind.database import SCHEMA_VERSION, Database, utcnow from imagefind.main import create_app from imagefind.media import ExtractedFrame from imagefind.offline_helper import offline_helper_script from PIL import Image def _api_app(tmp_path: Path, *, embedding_backend: str = "hash"): settings = Settings(data_dir=tmp_path / "data", embedding_backend=embedding_backend, upload_reserve_gb=0) settings.prepare() app = create_app(settings) _, token = app.state.services.auth.create_api_token("test") return app, {"Authorization": f"Bearer {token}"} def test_schema_v4_tracks_each_index_capability(tmp_path: Path): db = Database(tmp_path / "database.sqlite3") db.initialize() with db.read() as conn: columns = { row["name"] for row in conn.execute( "SELECT column_name AS name FROM information_schema.columns " "WHERE table_schema='public' AND table_name='videos'" ) } assert db.setting("schema_version") == SCHEMA_VERSION assert { "basic_fingerprint", "visual_model_version", "ocr_model_version", "faces_model_version", "audio_model_version", } <= columns def test_remember_device_cookie_defaults_to_ninety_days(tmp_path: Path): app, _ = _api_app(tmp_path) app.state.services.auth.setup("a sufficiently secure password") async def scenario(): transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: remembered = await client.post("/api/v1/auth/login", json={"password": "a sufficiently secure password"}) cookie = remembered.headers["set-cookie"] assert "Max-Age=7776000" in cookie assert "expires=" in cookie.lower() assert "Path=/" in cookie assert "HttpOnly" in cookie assert "SameSite=lax" in cookie session_only = await client.post( "/api/v1/auth/login", json={"password": "a sufficiently secure password", "remember_device": False}, ) cookie = session_only.headers["set-cookie"] assert "Max-Age" not in cookie assert "expires=" not in cookie.lower() asyncio.run(scenario()) def test_model_sources_mirror_config_and_offline_helper(tmp_path: Path): app, headers = _api_app(tmp_path) async def scenario(): transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: status = (await client.get("/api/v1/models", headers=headers)).json() assert status["sources"]["visual_image"].endswith("sentence-transformers/clip-ViT-B-32") assert "RapidOCR" in status["sources"]["ocr"] configured = await client.patch( "/api/v1/models/config", headers=headers, json={"hf_endpoint": "https://hf-mirror.example/base/"}, ) assert configured.json()["hf_endpoint"] == "https://hf-mirror.example/base" assert app.state.services.db.setting("model_hf_endpoint") == "https://hf-mirror.example/base" rejected = await client.patch( "/api/v1/models/config", headers=headers, json={"hf_endpoint": "https://user:secret@hf.example"}, ) assert rejected.status_code == 400 helper = await client.get("/api/v1/models/offline-helper", headers=headers) assert helper.status_code == 200 assert "prepare-imagefind-models.py" in helper.headers["content-disposition"] assert "--hf-endpoint" in helper.text assert "imagefind-models-" in helper.text compile(helper.text, "prepare-imagefind-models.py", "exec") asyncio.run(scenario()) def test_manifest_tamper_is_rejected_before_atomic_switch(tmp_path: Path): app, _ = _api_app(tmp_path) target = app.state.services.settings.models_dir (target / "visual" / "image").mkdir(parents=True) (target / "visual" / "text").mkdir(parents=True) sentinel = target / "visual" / "image" / "sentinel.bin" sentinel.write_bytes(b"current model") (target / "manifest.json").write_text(json.dumps({"version": "current"}), encoding="utf-8") files = { "visual/image/model.bin": b"tampered image", "visual/text/model.bin": b"text model", } manifest = { "format_version": 2, "version": "tampered", "source": "test", "files": { "visual/image/model.bin": "0" * 64, "visual/text/model.bin": hashlib.sha256(files["visual/text/model.bin"]).hexdigest(), }, } output = io.BytesIO() with tarfile.open(fileobj=output, mode="w:gz") as archive: for name, content in {**files, "manifest.json": json.dumps(manifest).encode()}.items(): member = tarfile.TarInfo(name) member.size = len(content) archive.addfile(member, io.BytesIO(content)) archive_path = tmp_path / "tampered.tar.gz" archive_path.write_bytes(output.getvalue()) try: app.state.services.models.install(archive_path=archive_path) except ValueError as exc: assert "篡改" in str(exc) else: raise AssertionError("tampered model package must be rejected") assert sentinel.read_bytes() == b"current model" assert json.loads((target / "manifest.json").read_text())["version"] == "current" def test_no_model_still_parses_and_text_searches_then_reconciles(tmp_path: Path): app, headers = _api_app(tmp_path, embedding_backend="auto") service = app.state.services now = utcnow() video_path = tmp_path / "Offline Holiday.mp4" video_path.write_bytes(b"video") with service.db.transaction() as conn: conn.execute( "INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)", ("source", "local", "本地", json.dumps({"path": str(tmp_path)}), now, now), ) conn.execute( "INSERT INTO videos(id,source_id,source_key,display_name,location,fingerprint,created_at,updated_at) " "VALUES(?,?,?,?,?,?,?,?)", ("video", "source", video_path.name, video_path.name, str(video_path), "fingerprint-v1", now, now), ) extraction = tmp_path / "extraction" extraction.mkdir() frame_path = extraction / "frame.jpg" Image.new("RGB", (320, 180), "navy").save(frame_path) service.media.input_for = lambda video: video_path service.media.probe = lambda media: { "raw": {"streams": []}, "duration_ms": 12_000, "width": 320, "height": 180, "codec": "h264", "container": "mp4", } service.media.extract_embedded_subtitles = lambda media, probe: [] service.media.sidecar_subtitles = lambda video: [] service.media.extract_frames = lambda media, duration: ( extraction, [ExtractedFrame(frame_path, 0, 0, 8_000)], ) service.indexer.index("test-job", "video") with service.db.read() as conn: video = conn.execute("SELECT * FROM videos WHERE id='video'").fetchone() assert video["status"] == "indexed" assert video["basic_fingerprint"] == "fingerprint-v1" assert video["visual_model_version"] is None assert video["ocr_model_version"] is None assert video["faces_model_version"] is None assert conn.execute("SELECT count(*) FROM frames WHERE video_id='video'").fetchone()[0] == 1 assert conn.execute("SELECT count(*) FROM text_entries WHERE video_id='video'").fetchone()[0] == 2 result = service.search.search( text="Offline Holiday", image_path=None, source_ids=None, min_duration_ms=None, max_duration_ms=None, min_width=None, has_people=None, tag_ids=None, sort="relevance", limit=10, cursor=None, ) assert result["items"][0]["video_id"] == "video" async def scenario(): transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: image_search = await client.post( "/api/v1/search", headers=headers, json={"image_id": "missing-query-image"} ) assert image_search.status_code == 409 assert "安装模型包" in image_search.json()["detail"] asyncio.run(scenario()) def test_search_cursor_keeps_a_stable_bounded_session(tmp_path: Path): app, _ = _api_app(tmp_path) search = app.state.services.search original = [{"frame_id": f"frame-{index}", "score": 1 / (index + 1)} for index in range(7)] first = search._store_page(original, 3) assert [item["frame_id"] for item in first["items"]] == ["frame-0", "frame-1", "frame-2"] assert first["total_candidates"] == 7 original[3]["frame_id"] = "mutated-after-cache" second = search._cached_page(first["next_cursor"], 3) assert [item["frame_id"] for item in second["items"]] == ["frame-3", "frame-4", "frame-5"] third = search._cached_page(second["next_cursor"], 3) assert [item["frame_id"] for item in third["items"]] == ["frame-6"] assert third["next_cursor"] is None session_id, _ = search._decode_cursor(first["next_cursor"]) with search._session_guard: search._sessions[session_id] = (time.monotonic() - 1, search._sessions[session_id][1]) try: search._cached_page(first["next_cursor"], 3) except ValueError as exc: assert "已过期" in str(exc) else: raise AssertionError("expired search cursor must not rerun a different result set") def test_hash_model_reconcile_queues_outdated_video(tmp_path: Path): app, headers = _api_app(tmp_path, embedding_backend="hash") service = app.state.services now = utcnow() with service.db.transaction() as conn: conn.execute( "INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)", ("source", "local", "本地", json.dumps({"path": str(tmp_path)}), now, now), ) conn.execute( "INSERT INTO videos(id,source_id,source_key,display_name,location,fingerprint,basic_fingerprint," "status,created_at,updated_at) VALUES(?,?,?,?,?,?,?,'indexed',?,?)", ("video", "source", "video.mp4", "video.mp4", str(tmp_path / "video.mp4"), "v1", "v1", now, now), ) async def scenario(): transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: status = (await client.get("/api/v1/models", headers=headers)).json() assert status["pending_videos"] == 1 response = await client.post("/api/v1/index/reconcile", headers=headers) assert response.status_code == 202 assert response.json()["queued"] == 1 asyncio.run(scenario()) def test_pending_videos_ignores_uninstalled_components(tmp_path: Path): app, headers = _api_app(tmp_path, embedding_backend="hash") service = app.state.services now = utcnow() with service.db.transaction() as conn: conn.execute( "INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)", ("source", "local", "本地", json.dumps({"path": str(tmp_path)}), now, now), ) conn.execute( "INSERT INTO videos(id,source_id,source_key,display_name,location,fingerprint,basic_fingerprint," "visual_model_version,status,available,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,'indexed',1,?,?)", ( "video", "source", "video.mp4", "video.mp4", str(tmp_path / "video.mp4"), "v1", "v1", "hash-v1", now, now, ), ) async def scenario(): transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: status = (await client.get("/api/v1/models", headers=headers)).json() assert status["visual_ready"] is True assert status["ocr_ready"] is False assert status["faces_ready"] is False assert status["audio_ready"] is False assert status["pending_videos"] == 0 asyncio.run(scenario()) def test_offline_helper_source_is_standalone_python(): source = offline_helper_script() assert "snapshot_download" in source assert "Range" in source compile(source, "offline-helper.py", "exec")