imagefind/tests/test_offline_models_and_index_v4.py

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")