feat: add ImageFind application and release pipelines

This commit is contained in:
2026-08-11 18:02:40 +08:00
commit 16239d7525
270 changed files with 59163 additions and 0 deletions
+93
View File
@@ -0,0 +1,93 @@
from __future__ import annotations
import os
import sys
import threading
import uuid
from pathlib import Path
import pytest
from imagefind import api as api_module
from imagefind.database import Database
@pytest.fixture(autouse=True)
def _isolated_postgres_database(monkeypatch):
"""Give every distinct test data directory its own PostgreSQL database.
Set ``IMAGEFIND_TEST_POSTGRES_ADMIN_DSN`` to a PostgreSQL/pgvector server
where the configured user may create and drop databases. Tests which do
not open ``Database`` remain usable without that environment variable.
"""
admin_dsn = os.environ.get("IMAGEFIND_TEST_POSTGRES_ADMIN_DSN", "").strip()
if not admin_dsn:
yield
return
import psycopg
from psycopg import sql
from psycopg.conninfo import conninfo_to_dict
admin_parameters = conninfo_to_dict(admin_dsn)
admin_parameters.setdefault("dbname", "postgres")
configurations: dict[Path, tuple[str, Path]] = {}
guard = threading.Lock()
def postgres_conf(database: Database) -> Path:
data_dir = database.path.parent.resolve()
with guard:
existing = configurations.get(data_dir)
if existing:
return existing[1]
database_name = f"imagefind_test_{uuid.uuid4().hex[:24]}"
with psycopg.connect(**admin_parameters, autocommit=True) as admin:
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(database_name)))
conf_path = data_dir / ".postgres-client.conf"
conf_path.parent.mkdir(parents=True, exist_ok=True)
conf_path.write_text(
"\n".join(
(
f"host={admin_parameters.get('host', '127.0.0.1')}",
f"port={admin_parameters.get('port', '5432')}",
f"database={database_name}",
f"username={admin_parameters.get('user', 'postgres')}",
f"password={admin_parameters.get('password', '')}",
f"sslmode={admin_parameters.get('sslmode', 'disable')}",
"",
)
),
encoding="utf-8",
)
configurations[data_dir] = (database_name, conf_path)
return conf_path
monkeypatch.setattr(Database, "postgres_conf_path", property(postgres_conf))
try:
yield
finally:
for database_name, _ in configurations.values():
with psycopg.connect(**admin_parameters, autocommit=True) as admin:
admin.execute(
sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(
sql.Identifier(database_name)
)
)
@pytest.fixture(autouse=True)
def _python_313_asgi_thread_compatibility(monkeypatch):
"""Avoid a Python 3.13 sandbox-only ASGI selector deadlock.
Production fnOS uses Python 3.12 and keeps blocking API reads in worker
threads. The repository test sandbox can deadlock when ASGITransport awaits
``asyncio.to_thread`` on 3.13, so tests execute only this wrapper inline.
"""
if sys.version_info < (3, 13):
return
async def inline(function, /, *args, **kwargs):
return function(*args, **kwargs)
monkeypatch.setattr(api_module, "_background_api", inline)
+726
View File
@@ -0,0 +1,726 @@
from __future__ import annotations
import os
import sys
import types
from pathlib import Path
import imagefind.accelerator as accelerator_module
import imagefind.ai as ai_module
import imagefind.main as main_module
import pytest
from imagefind.accelerator import AcceleratorService
from imagefind.ai import EmbeddingService, ModelUnavailable, _OpenVINOClipImageEncoder, _OpenVINOOrtSession
from imagefind.config import Settings
from imagefind.media import MediaService
def _openvino(monkeypatch, devices: list[str]):
class Core:
available_devices = devices
@staticmethod
def get_property(device, _name):
return {"CPU": "Intel CPU", "GPU.0": "Intel UHD Graphics"}.get(device, device)
monkeypatch.setitem(sys.modules, "openvino", types.SimpleNamespace(Core=Core))
def _render_nodes(monkeypatch, nodes: list[Path], accessible: bool = True):
original_glob = Path.glob
original_access = os.access
def fake_glob(path: Path, pattern: str):
if path == Path("/dev/dri") and pattern == "renderD*":
return iter(nodes)
return original_glob(path, pattern)
monkeypatch.setattr(Path, "glob", fake_glob)
monkeypatch.setattr(
accelerator_module.os,
"access",
lambda path, mode: accessible if str(path).startswith("/dev/dri/renderD") else original_access(path, mode),
)
def test_accelerator_prefers_gpu_and_falls_back_per_component(monkeypatch, tmp_path: Path):
_openvino(monkeypatch, ["CPU", "GPU.0"])
_render_nodes(monkeypatch, [Path("/dev/dri/renderD128")])
service = AcceleratorService(Settings(data_dir=tmp_path, ai_cpu_threads=3))
assert service.gpu_device == "GPU.0"
assert service.device_for("visual") == "GPU.0"
assert service.status()["device_names"]["GPU.0"] == "Intel UHD Graphics"
service.mark_ready("visual", "GPU.0")
assert service.fall_back("visual", RuntimeError("secret driver detail")) is True
assert service.device_for("visual") == "CPU"
assert service.device_for("ocr") == "GPU.0"
assert service.fall_back("visual", RuntimeError("do not retry")) is False
visual = service.status()["components"]["visual"]
assert visual["state"] == "fallback"
assert "secret driver detail" not in visual["fallback_reason"]
assert visual["requested_device"] == "GPU.0"
assert visual["actual_device"] == "CPU"
assert visual["failure_stage"] == "inference"
assert service.ov_config("CPU") == {"INFERENCE_NUM_THREADS": 3, "PERFORMANCE_HINT": "LATENCY"}
gpu_config = service.ov_config("GPU.0")
assert gpu_config["PERFORMANCE_HINT"] == "LATENCY"
assert gpu_config["NUM_STREAMS"] == "1"
assert gpu_config["CACHE_DIR"].endswith("openvino-cache")
service.reset("visual")
visual = service.status()["components"]["visual"]
assert visual["state"] == "not_loaded"
assert visual["device"] is None
assert visual["requested_device"] == "GPU.0"
assert visual["actual_device"] == "CPU"
assert visual["execution_devices"] == ["CPU"]
assert visual["last_verified_at"] is not None
assert visual["failure_stage"] is None
assert visual["fallback_reason"] is None
assert service.device_for("visual") == "GPU.0"
@pytest.mark.parametrize(
("mode", "hint", "streams", "batch_size"),
[("low", "LATENCY", "1", 1), ("balanced", "THROUGHPUT", "2", 2), ("throughput", "THROUGHPUT", "3", 4)],
)
def test_gpu_profiles_bound_streams_and_batches(monkeypatch, tmp_path: Path, mode, hint, streams, batch_size):
_openvino(monkeypatch, ["CPU", "GPU.0"])
_render_nodes(monkeypatch, [Path("/dev/dri/renderD128")])
service = AcceleratorService(Settings(data_dir=tmp_path, ai_gpu_mode=mode))
config = service.ov_config("GPU.0")
assert config["PERFORMANCE_HINT"] == hint
assert config["NUM_STREAMS"] == streams
assert service.gpu_profile()["batch_size"] == batch_size
assert service.settings.ai_cpu_threads == 1
def test_accelerator_reset_without_prior_inference_remains_never_loaded(monkeypatch, tmp_path: Path):
_openvino(monkeypatch, ["CPU", "GPU.0"])
_render_nodes(monkeypatch, [Path("/dev/dri/renderD128")])
service = AcceleratorService(Settings(data_dir=tmp_path))
service.reset("faces")
faces = service.status()["components"]["faces"]
assert faces["state"] == "not_loaded"
assert faces["requested_device"] is None
assert faces["actual_device"] is None
assert faces["execution_devices"] == []
assert faces["last_verified_at"] is None
def test_accelerator_cpu_verification_is_transient_but_can_become_fallback(
monkeypatch,
tmp_path: Path,
):
_openvino(monkeypatch, ["CPU", "GPU.0"])
_render_nodes(monkeypatch, [Path("/dev/dri/renderD128")])
service = AcceleratorService(Settings(data_dir=tmp_path))
service.mark_ready("audio", "GPU.0", ["GPU.0"])
service.mark_verifying_cpu("audio")
verifying = service.status()["components"]["audio"]
assert verifying["state"] == "verifying_cpu"
assert verifying["actual_device"] == "CPU"
assert service.device_for("audio") == "GPU.0"
assert service.fall_back("audio", "CPU sample contained speech", stage="empty_result") is True
fallback = service.status()["components"]["audio"]
assert fallback["state"] == "fallback"
assert fallback["actual_device"] == "CPU"
assert service.device_for("audio") == "CPU"
def test_audio_transient_device_failures_only_open_circuit_after_three_tasks(
monkeypatch,
tmp_path: Path,
):
_openvino(monkeypatch, ["CPU", "GPU.0"])
_render_nodes(monkeypatch, [Path("/dev/dri/renderD128")])
service = AcceleratorService(Settings(data_dir=tmp_path))
assert service.record_transient_failure("audio", "stall one", stage="inference_stall") is False
assert service.device_for("audio") == "GPU.0"
assert service.record_transient_failure("audio", "stall two", stage="inference_stall") is False
assert service.device_for("audio") == "GPU.0"
assert service.record_transient_failure("audio", "stall three", stage="inference_stall") is True
status = service.status()["components"]["audio"]
assert status["state"] == "fallback"
assert status["fallback_scope"] == "component"
assert status["circuit_state"] == "open"
assert status["failure_count"] == 3
assert status["retry_at"]
assert service.device_for("audio") == "CPU"
service.mark_inference_success("audio", "GPU.0", ["GPU.0"])
status = service.status()["components"]["audio"]
assert status["circuit_state"] == "closed"
assert status["failure_count"] == 0
def test_accelerator_reports_render_permission_problem(monkeypatch, tmp_path: Path):
_openvino(monkeypatch, ["CPU"])
_render_nodes(monkeypatch, [Path("/dev/dri/renderD128")], accessible=False)
service = AcceleratorService(Settings(data_dir=tmp_path))
assert service.gpu_device is None
assert service.status()["render_device"] == {"available": True, "accessible": False, "count": 1}
assert service.status()["unavailable_reason"] == "Intel render 设备权限不足"
def test_gpu_sysfs_metrics_need_two_samples_and_do_not_invent_shared_vram(
monkeypatch,
tmp_path: Path,
):
_openvino(monkeypatch, ["CPU", "GPU.0"])
_render_nodes(monkeypatch, [Path("/dev/dri/renderD128")])
service = AcceleratorService(Settings(data_dir=tmp_path))
busy_path = Path("/sys/class/drm/card0/engine/rcs0/busy")
original_glob = Path.glob
values = iter((1_000_000_000, 1_500_000_000))
def fake_glob(path: Path, pattern: str):
if path == Path("/sys/class/drm") and pattern == "card*/engine/*/busy":
return iter((busy_path,))
if path == Path("/sys/class/drm") and pattern == "card*/device":
return iter(())
return original_glob(path, pattern)
original_read_text = Path.read_text
monkeypatch.setattr(Path, "glob", fake_glob)
monkeypatch.setattr(
Path,
"read_text",
lambda path, **kwargs: str(next(values))
if path == busy_path
else original_read_text(path, **kwargs),
)
clock = iter((10.0, 11.0))
monkeypatch.setattr(accelerator_module.time, "monotonic", lambda: next(clock))
first = service.hardware_metrics()
second = service.hardware_metrics()
assert first["supported"] is True
assert first["utilization_percent"] is None
assert second["utilization_percent"] == 50.0
assert second["engines"] == {"card0/rcs0": 50.0}
assert second["memory_supported"] is False
assert second["memory_total_bytes"] is None
def test_visual_models_refuse_runtime_export_and_require_persisted_ir(tmp_path: Path):
settings = Settings(data_dir=tmp_path)
for name in ("image", "text"):
root = settings.models_dir / "visual" / name
root.mkdir(parents=True)
(root / "modules.json").write_text('[{"idx":0,"path":"","type":"Transformer"}]')
class Accelerator:
@staticmethod
def ov_config(_device):
return {"PERFORMANCE_HINT": "THROUGHPUT"}
@staticmethod
def mark_ready(_component, _device):
pass
calls = []
class Model:
def __init__(self, path, **kwargs):
calls.append((path, kwargs))
service = EmbeddingService(settings, Accelerator())
with pytest.raises(ModelUnavailable, match="OpenVINO 模型不完整"):
service._build_models(Model, "openvino", "GPU.0")
assert calls == []
@pytest.mark.parametrize("text_module", ["", "0_Transformer"])
def test_visual_models_reuse_existing_openvino_exports(
tmp_path: Path, monkeypatch, text_module: str
):
settings = Settings(data_dir=tmp_path)
for name in ("image", "text"):
root = settings.models_dir / "visual" / name
module = root / text_module if name == "text" and text_module else root
export = module / "openvino"
export.mkdir(parents=True)
module_type = "Transformer" if name == "text" else "CLIPModel"
module_path = text_module if name == "text" else ""
(root / "modules.json").write_text(
f'['
f'{{"idx":0,"path":"{module_path}","type":"{module_type}"}}'
f']'
)
(module / "config.json").write_text("{}")
(export / "openvino_model.xml").write_text("<xml/>")
(export / "openvino_model.bin").write_bytes(b"model")
(export / "config.json").write_text("{}")
class Accelerator:
@staticmethod
def ov_config(_device):
return {"PERFORMANCE_HINT": "THROUGHPUT"}
@staticmethod
def execution_devices_from(*_models):
return ["GPU.0"]
@staticmethod
def mark_ready(_component, _device, _execution_devices):
pass
text_calls = []
image_calls = []
class ImageModel:
def __init__(self, path, device, ov_config):
image_calls.append((path, device, ov_config))
class Model:
def __init__(self, path, **kwargs):
text_calls.append((path, kwargs))
monkeypatch.setattr(ai_module, "_OpenVINOClipImageEncoder", ImageModel)
EmbeddingService(settings, Accelerator())._build_models(Model, "openvino", "GPU.0")
assert len(image_calls) == 1
assert Path(image_calls[0][0]).is_absolute()
assert image_calls[0][1:] == ("GPU", {"PERFORMANCE_HINT": "THROUGHPUT"})
assert len(text_calls) == 1
assert Path(text_calls[0][0]).is_absolute()
assert text_calls[0][1]["local_files_only"] is True
assert text_calls[0][1]["model_kwargs"] == {
"device": "GPU",
"ov_config": {"PERFORMANCE_HINT": "THROUGHPUT"},
"export": False,
"file_name": "openvino_model.xml",
"subfolder": "openvino",
}
def test_clip_image_encoder_uses_image_aware_openvino_model(monkeypatch, tmp_path: Path):
root = tmp_path / "image"
export = root / "openvino"
export.mkdir(parents=True)
processor_root = root / "0_CLIPModel"
processor_root.mkdir()
(root / "modules.json").write_text(
'[{"idx":0,"path":"0_CLIPModel","type":"sentence_transformers.models.CLIPModel"}]'
)
(processor_root / "config.json").write_text("{}")
(export / "openvino_model.xml").write_text("<xml/>")
(export / "openvino_model.bin").write_bytes(b"model")
(export / "config.json").write_text("{}")
calls = {}
class Matrix:
def __init__(self, rows):
self.rows = rows
def __truediv__(self, denominators):
return Matrix(
[[value / denominators[index][0] for value in row] for index, row in enumerate(self.rows)]
)
def tolist(self):
return self.rows
numpy = types.ModuleType("numpy")
numpy.float32 = "float32"
numpy.asarray = lambda value, dtype=None: Matrix(value)
numpy.maximum = lambda values, _minimum: values
numpy.finfo = lambda _dtype: types.SimpleNamespace(eps=1e-7)
numpy.linalg = types.SimpleNamespace(
norm=lambda values, axis, keepdims: [
[sum(item * item for item in row) ** 0.5] for row in values.rows
]
)
class OpenVINOModel:
@classmethod
def from_pretrained(cls, _path, **_kwargs):
raise AssertionError("public loader must not perform Hugging Face library inference")
@classmethod
def _from_pretrained(cls, path, **kwargs):
calls["model"] = (path, kwargs)
return cls()
def __call__(self, **inputs):
calls["inputs"] = inputs
return types.SimpleNamespace(image_embeds=[[3.0, 4.0]])
class Processor:
@classmethod
def from_pretrained(cls, path, **kwargs):
calls["processor"] = (path, kwargs)
return cls()
def __call__(self, **kwargs):
calls["processor_inputs"] = kwargs
return {"input_ids": [[1]], "pixel_values": [[0.0]]}
class Config:
@classmethod
def from_pretrained(cls, path, **kwargs):
calls["config"] = (path, kwargs)
return cls()
optimum = types.ModuleType("optimum")
optimum_intel = types.ModuleType("optimum.intel")
optimum_openvino = types.ModuleType("optimum.intel.openvino")
optimum_openvino.OVModelForZeroShotImageClassification = OpenVINOModel
transformers = types.ModuleType("transformers")
transformers.AutoConfig = Config
transformers.AutoProcessor = Processor
for name, module in {
"optimum": optimum,
"optimum.intel": optimum_intel,
"optimum.intel.openvino": optimum_openvino,
"transformers": transformers,
"numpy": numpy,
}.items():
monkeypatch.setitem(sys.modules, name, module)
encoder = _OpenVINOClipImageEncoder(root, "GPU.0", {"NUM_STREAMS": "1"})
vectors = encoder.encode([object()])
assert calls["model"] == (
str(export.resolve()),
{
"config": calls["model"][1]["config"],
"file_name": "openvino_model.xml",
"local_files_only": True,
"device": "GPU.0",
"ov_config": {"NUM_STREAMS": "1"},
},
)
assert isinstance(calls["model"][1]["config"], Config)
assert calls["config"] == (str(processor_root.resolve()), {"local_files_only": True})
assert calls["processor"] == (str(processor_root.resolve()), {"local_files_only": True})
assert calls["processor_inputs"]["text"] == [""]
assert set(calls["inputs"]) == {"input_ids", "pixel_values"}
assert [round(value, 6) for value in vectors.tolist()[0]] == [0.6, 0.8]
def test_visual_text_encoder_bypasses_optimum_hub_inference(
tmp_path: Path, monkeypatch
):
settings = Settings(data_dir=tmp_path)
for name in ("image", "text"):
root = settings.models_dir / "visual" / name
export = root / "openvino"
export.mkdir(parents=True)
module_type = "CLIPModel" if name == "image" else "Transformer"
(root / "modules.json").write_text(
f'[' f'{{"idx":0,"path":"","type":"{module_type}"}}' f']'
)
(root / "config.json").write_text("{}")
(export / "openvino_model.xml").write_text("<xml/>")
(export / "openvino_model.bin").write_bytes(b"model")
(export / "config.json").write_text("{}")
calls = {}
class OpenVINOTextModel:
@classmethod
def from_pretrained(cls, _path, **_kwargs):
raise AssertionError("public Optimum loader must not infer a Hub library")
@classmethod
def _from_pretrained(cls, **kwargs):
calls["local"] = kwargs
return object()
class SentenceTransformerModel:
__module__ = "sentence_transformers.sentence_transformer.model"
def __init__(self, path, **kwargs):
calls["wrapper"] = (path, kwargs)
options = kwargs["model_kwargs"]
self.model = OpenVINOTextModel.from_pretrained(
path,
config="local-config",
export=options["export"],
file_name=options["file_name"],
subfolder=options["subfolder"],
local_files_only=kwargs["local_files_only"],
device=options["device"],
ov_config=options["ov_config"],
)
optimum = types.ModuleType("optimum")
optimum_intel = types.ModuleType("optimum.intel")
optimum_openvino = types.ModuleType("optimum.intel.openvino")
optimum_openvino.OVModelForFeatureExtraction = OpenVINOTextModel
for name, module in {
"optimum": optimum,
"optimum.intel": optimum_intel,
"optimum.intel.openvino": optimum_openvino,
}.items():
monkeypatch.setitem(sys.modules, name, module)
class Accelerator:
@staticmethod
def ov_config(_device):
return {"NUM_STREAMS": "1"}
@staticmethod
def execution_devices_from(*_models):
return ["GPU"]
@staticmethod
def mark_ready(*_args):
pass
monkeypatch.setattr(ai_module, "_OpenVINOClipImageEncoder", lambda *_args: object())
EmbeddingService(settings, Accelerator())._build_models(
SentenceTransformerModel, "openvino", "GPU"
)
assert calls["local"] == {
"model_id": str((settings.models_dir / "visual" / "text").resolve()),
"config": "local-config",
"file_name": "openvino_model.xml",
"subfolder": "openvino",
"local_files_only": True,
"device": "GPU",
"ov_config": {"NUM_STREAMS": "1"},
}
with pytest.raises(AssertionError, match="public Optimum loader"):
OpenVINOTextModel.from_pretrained("/vol1/imagefind/models/visual/text")
def test_hugging_face_validation_errors_are_not_gpu_device_errors():
error_type = type("HFValidationError", (ValueError,), {"__module__": "huggingface_hub.errors"})
assert AcceleratorService.is_device_error(error_type("invalid repo id")) is False
assert AcceleratorService.is_device_error(ModelUnavailable("OpenVINO IR 尚未准备完成")) is False
assert AcceleratorService.is_device_error(RuntimeError("OpenVINO GPU compile_model failed")) is True
def test_execution_devices_are_read_from_compiled_wrappers():
class Compiled:
@staticmethod
def get_property(name):
assert name == "EXECUTION_DEVICES"
return ["GPU.0"]
class Request:
@staticmethod
def get_compiled_model():
return Compiled()
wrapper = types.SimpleNamespace(encoder=types.SimpleNamespace(request=Request()))
assert AcceleratorService.execution_devices_from(wrapper) == ["GPU.0"]
def test_execution_device_discovery_ignores_unsupported_gpu_properties():
class UnsupportedProperties:
@staticmethod
def get_property(_name):
raise ValueError("property is not supported by this GPU wrapper")
@staticmethod
def get_compiled_model():
raise ValueError("compiled model is exposed by a sibling request")
class Compiled:
@staticmethod
def get_property(name):
assert name == "EXECUTION_DEVICES"
return ["GPU.0"]
wrapper = types.SimpleNamespace(
model=UnsupportedProperties(),
request=Compiled(),
)
assert AcceleratorService.execution_devices_from(wrapper) == ["GPU.0"]
def test_openvino_value_errors_retry_on_cpu_but_model_layout_errors_do_not():
assert AcceleratorService.is_device_error(ValueError("GPU plugin property failed")) is True
assert AcceleratorService.is_device_error(ValueError("modules.json is invalid")) is False
def test_visual_gpu_value_error_retries_complete_model_on_cpu(tmp_path: Path, monkeypatch):
settings = Settings(data_dir=tmp_path)
for name in ("image", "text"):
root = settings.models_dir / "visual" / name
export = root / "openvino"
export.mkdir(parents=True)
(root / "modules.json").write_text('[{"idx":0,"path":"","type":"Transformer"}]')
(root / "config.json").write_text("{}")
(export / "openvino_model.xml").write_text("<xml/>")
(export / "openvino_model.bin").write_bytes(b"model")
(export / "config.json").write_text("{}")
calls = []
unavailable = []
class Accelerator:
def __init__(self):
self.fallback = False
def device_for(self, _component):
return "CPU" if self.fallback else "GPU.0"
@staticmethod
def is_device_error(error):
return isinstance(error, ValueError)
def fall_back(self, _component, _error, stage):
assert stage == "image_encoder_compile"
self.fallback = True
return True
@staticmethod
def ov_config(device):
return {"device_config": device}
@staticmethod
def execution_devices_from(*_models):
return ["CPU"]
@staticmethod
def mark_ready(component, device, execution_devices):
calls.append((component, device, execution_devices))
@staticmethod
def mark_unavailable(component, error, stage):
unavailable.append((component, error, stage))
class ImageModel:
def __init__(self, _path, device, _ov_config):
if device == "GPU":
raise ValueError("Intel GPU property rejected by wrapper")
class TextModel:
def __init__(self, _path, **kwargs):
assert kwargs["model_kwargs"]["device"] == "CPU"
sentence_transformers = types.ModuleType("sentence_transformers")
sentence_transformers.SentenceTransformer = TextModel
monkeypatch.setitem(sys.modules, "sentence_transformers", sentence_transformers)
monkeypatch.setattr(ai_module, "_OpenVINOClipImageEncoder", ImageModel)
service = EmbeddingService(settings, Accelerator())
service._load()
assert service._image_model is not None
assert service._text_model is not None
assert calls == [("visual", "CPU", ["CPU"])]
assert unavailable == []
def test_openvino_ocr_adapter_compiles_requested_device(monkeypatch):
captured = {}
class Metadata:
value = "a\nb"
class Model:
@staticmethod
def get_rt_info():
return {"framework": {"character": Metadata()}}
class Compiled:
@staticmethod
def output(_index):
return "output"
def __call__(self, values):
captured["values"] = values
return {"output": "result"}
class Core:
@staticmethod
def read_model(path):
captured["model"] = path
return Model()
@staticmethod
def compile_model(model, device, config):
captured.update(compiled_model=model, device=device, config=config)
return Compiled()
monkeypatch.setitem(sys.modules, "openvino", types.SimpleNamespace(Core=Core))
session = _OpenVINOOrtSession("ocr.onnx", "GPU.0", {"PERFORMANCE_HINT": "THROUGHPUT"})
assert session("pixels") == ["result"]
assert session.have_key() is True
assert session.get_character_list() == ["a", "b"]
assert captured["model"] == "ocr.onnx"
assert isinstance(captured["compiled_model"], Model)
assert captured["device"] == "GPU.0"
assert captured["config"] == {"PERFORMANCE_HINT": "THROUGHPUT"}
assert captured["values"] == ["pixels"]
def test_remote_media_uses_runtime_internal_port(tmp_path: Path):
settings = Settings(data_dir=tmp_path, port=8765, internal_media_port=43123)
media = MediaService(settings).input_for({"id": "remote-video", "source_kind": "webdav"})
assert media.value == "http://127.0.0.1:43123/api/internal/remote/remote-video"
assert media.headers and media.headers.startswith("X-ImageFind-Internal: ")
assert media.headers.endswith("\r\n")
def test_gateway_only_serve_creates_private_random_tcp_listener(monkeypatch, tmp_path: Path):
gateway_socket = tmp_path / "imagefind.sock"
settings = Settings(
data_dir=tmp_path / "data",
gateway_socket=gateway_socket,
direct_access=False,
host="0.0.0.0",
port=8765,
internal_media_port=0,
)
captured = {}
class Listener:
def __init__(self, name):
self.name = name
def getsockname(self):
return self.name
def close(self):
pass
class Config:
def __init__(self, app, **kwargs):
captured.update(app=app, config=kwargs)
class Server:
def __init__(self, _config):
pass
@staticmethod
def run(sockets):
captured["listeners"] = [listener.getsockname() for listener in sockets]
monkeypatch.setattr(main_module, "create_app", lambda current: ("app", current.internal_media_port))
monkeypatch.setattr(
main_module,
"_tcp_socket",
lambda host, port: captured.setdefault("tcp_request", (host, port)) and Listener((host, 43123)),
)
monkeypatch.setattr(main_module, "_unix_socket", lambda path, _mode: Listener(str(path)))
monkeypatch.setitem(sys.modules, "uvicorn", types.SimpleNamespace(Config=Config, Server=Server))
main_module.serve(settings)
assert captured["tcp_request"] == ("127.0.0.1", 0)
assert settings.internal_media_port == 43123
assert captured["app"] == ("app", settings.internal_media_port)
assert captured["listeners"][0] == ("127.0.0.1", settings.internal_media_port)
assert captured["listeners"][1] == str(gateway_socket)
assert all(listener != ("0.0.0.0", 8765) for listener in captured["listeners"])
assert not gateway_socket.exists()
+842
View File
@@ -0,0 +1,842 @@
import asyncio
import base64
import json
from pathlib import Path
import httpx
import imagefind.api as api_module
from imagefind import __version__
from imagefind.config import Settings
from imagefind.database import utcnow
from imagefind.main import create_app
from imagefind.speech import SPEECH_INDEX_REVISION
from imagefind.text import search_tokens
def test_setup_login_csrf_and_token(tmp_path: Path):
app = create_app(Settings(data_dir=tmp_path, embedding_backend="hash", scan_interval_seconds=86400))
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
assert (await client.get("/api/v1/status")).json()["configured"] is False
response = await client.post("/api/v1/setup", json={"password": "a secure test password"})
assert response.status_code == 200
csrf = response.json()["csrf_token"]
assert (await client.get("/api/v1/auth/me")).status_code == 200
assert (await client.post("/api/v1/tokens", json={"name": "test"})).status_code == 403
token = await client.post(
"/api/v1/tokens",
json={"name": "test"},
headers={"X-CSRF-Token": csrf},
)
assert token.status_code == 201
api_token = token.json()["token"]
assert (
await client.get("/api/v1/sources", headers={"Authorization": f"Bearer {api_token}"})
).status_code == 200
asyncio.run(scenario())
def test_status_frontend_fallback_and_openapi(tmp_path: Path):
settings = Settings(
data_dir=tmp_path,
frontend_dir=tmp_path / "frontend-not-built",
embedding_backend="hash",
scan_interval_seconds=86400,
)
app = create_app(settings)
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/status")
assert status.status_code == 200
assert status.json() == {"configured": False, "version": __version__, "access_mode": "direct"}
root = await client.get("/")
assert root.status_code == 200
assert "ImageFind API 正在运行" in root.text
openapi = await client.get("/api/openapi.json")
assert openapi.status_code == 200
assert "/api/v1/search" in openapi.json()["paths"]
asyncio.run(scenario())
def test_security_headers_token_scope_payload_and_diagnostics(tmp_path: Path):
app = create_app(Settings(data_dir=tmp_path, embedding_backend="hash", scan_interval_seconds=86400))
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/status")
assert status.headers["x-content-type-options"] == "nosniff"
assert status.headers["referrer-policy"] == "same-origin"
assert "camera=()" in status.headers["permissions-policy"]
assert "default-src 'self'" in status.headers["content-security-policy"]
setup = await client.post("/api/v1/setup", json={"password": "a secure test password"})
csrf = setup.json()["csrf_token"]
created = await client.post(
"/api/v1/tokens",
headers={"X-CSRF-Token": csrf},
json={"name": "DAV automation", "scopes": ["webdav"]},
)
assert created.status_code == 201
assert created.json()["scopes"] == ["webdav"]
listed = await client.get("/api/v1/tokens")
assert listed.json()[0]["scopes"] == ["webdav"]
assert "token" not in listed.json()[0]
diagnostics = await client.get("/api/v1/system/diagnostics")
assert diagnostics.status_code == 200
payload = diagnostics.json()
assert payload["process_rss_bytes"] >= 0
assert payload["database"]["engine"] == "postgresql"
assert payload["database"]["pool_max"] >= 1
assert isinstance(payload["database"]["activity"]["states"], dict)
assert "running" in payload["inference"]
assert payload["events"]["subscribers"] == 0
asyncio.run(scenario())
def test_fnos_gateway_auth_is_bound_to_unix_socket_and_prefix(tmp_path: Path):
socket_path = tmp_path / "imagefind.sock"
settings = Settings(
data_dir=tmp_path / "data",
gateway_socket=socket_path,
gateway_prefix="/app/imagefind",
embedding_backend="hash",
scan_interval_seconds=86400,
)
app = create_app(settings)
app.state.services.auth.setup("gateway fallback password")
class UnixSocketScope:
def __init__(self, target):
self.target = target
async def __call__(self, scope, receive, send):
mounted = dict(scope)
mounted["server"] = (str(socket_path), None)
await self.target(mounted, receive, send)
async def scenario():
# Header spoofing over TCP must never enable gateway SSO.
direct_transport = httpx.ASGITransport(app=app)
spoofed = {
"X-Trim-Isadmin": "true",
"X-Trim-Userid": "nas-admin",
"X-Trim-Username": "Administrator",
}
async with httpx.AsyncClient(transport=direct_transport, base_url="http://test") as direct:
status = await direct.get("/api/v1/status", headers=spoofed)
assert status.json()["access_mode"] == "direct"
assert (await direct.post("/api/v1/auth/gateway", headers=spoofed)).status_code == 404
assert (await direct.get("/api/v1/sources", headers=spoofed)).status_code == 401
gateway_transport = httpx.ASGITransport(app=UnixSocketScope(app))
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers=spoofed,
follow_redirects=False,
) as gateway:
redirect = await gateway.get("/app/imagefind")
assert redirect.status_code == 307
assert redirect.headers["location"] == "/app/imagefind/"
assert (await gateway.get("/outside-prefix")).status_code == 404
status = await gateway.get("/app/imagefind/api/v1/status")
assert status.json()["access_mode"] == "gateway"
login = await gateway.post("/app/imagefind/api/v1/auth/gateway")
assert login.status_code == 200
csrf = login.json()["csrf_token"]
gateway_token = login.json()["gateway_session_token"]
assert gateway_token
assert login.json()["nas_username"] == "Administrator"
cookie = login.headers["set-cookie"]
assert "imagefind_gateway_session=" in cookie
assert "Path=/app/imagefind/" in cookie
me = await gateway.get("/app/imagefind/api/v1/auth/me")
assert me.status_code == 200
assert me.json()["kind"] == "gateway"
assert (await gateway.post("/app/imagefind/api/v1/tokens", json={"name": "blocked"})).status_code == 403
created = await gateway.post(
"/app/imagefind/api/v1/tokens",
json={"name": "gateway"},
headers={**spoofed, "X-CSRF-Token": csrf},
)
assert created.status_code == 201
changed_user = await gateway.get(
"/app/imagefind/api/v1/auth/me",
headers={**spoofed, "X-Trim-Userid": "another-admin"},
)
assert changed_user.status_code == 401
# The fnOS WebView may omit the scoped cookie. The short-lived session
# header keeps SSO working, but only inside the trusted Unix gateway.
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers={**spoofed, "X-ImageFind-Gateway-Session": gateway_token},
) as header_only:
me = await header_only.get("/app/imagefind/api/v1/auth/me")
assert me.status_code == 200
assert me.json()["nas_user_id"] == "nas-admin"
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers={
**spoofed,
"X-Trim-Userid": "another-admin",
"X-ImageFind-Gateway-Session": gateway_token,
},
) as wrong_identity:
assert (await wrong_identity.get("/app/imagefind/api/v1/auth/me")).status_code == 401
other_headers = {
"X-Trim-Isadmin": "true",
"X-Trim-Userid": "other-admin",
"X-Trim-Username": "Other administrator",
}
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers=other_headers,
) as other_gateway:
other_login = await other_gateway.post("/app/imagefind/api/v1/auth/gateway")
other_token = other_login.json()["gateway_session_token"]
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers={
**spoofed,
"Cookie": f"imagefind_gateway_session={other_token}",
"X-ImageFind-Gateway-Session": gateway_token,
},
) as stale_cookie:
me = await stale_cookie.get("/app/imagefind/api/v1/auth/me")
assert me.status_code == 200
assert me.json()["nas_user_id"] == "nas-admin"
async with httpx.AsyncClient(transport=direct_transport, base_url="http://test") as direct:
response = await direct.get(
"/api/v1/auth/me",
headers={**spoofed, "X-ImageFind-Gateway-Session": gateway_token},
)
assert response.status_code == 401
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers={"X-Trim-Isadmin": "false", "X-Trim-Userid": "ordinary-user"},
) as ordinary:
assert (await ordinary.post("/app/imagefind/api/v1/auth/gateway")).status_code == 403
async with httpx.AsyncClient(transport=gateway_transport, base_url="http://test") as anonymous:
assert (await anonymous.post("/app/imagefind/api/v1/auth/gateway")).status_code == 401
asyncio.run(scenario())
def test_frontend_runtime_base_is_injected_for_gateway_and_direct_access(tmp_path: Path):
frontend = tmp_path / "frontend"
frontend.mkdir()
(frontend / "index.html").write_text("<html><head></head><body>ImageFind</body></html>")
socket_path = tmp_path / "imagefind.sock"
settings = Settings(
data_dir=tmp_path / "data",
frontend_dir=frontend,
gateway_socket=socket_path,
embedding_backend="hash",
scan_interval_seconds=86400,
)
app = create_app(settings)
class UnixSocketScope:
async def __call__(self, scope, receive, send):
mounted = dict(scope)
mounted["server"] = (str(socket_path), None)
await app(mounted, receive, send)
async def scenario():
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as direct:
index = await direct.get("/")
assert 'window.__IMAGEFIND_BASE__="/"' in index.text
assert 'dataset.imagefindAccess="direct"' in index.text
assert index.headers["cache-control"] == "no-store"
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=UnixSocketScope()), base_url="http://test"
) as gateway:
index = await gateway.get("/app/imagefind/")
assert 'window.__IMAGEFIND_BASE__="/app/imagefind/"' in index.text
assert 'dataset.imagefindAccess="gateway"' in index.text
asyncio.run(scenario())
def _seed_local_search(app, tmp_path: Path) -> tuple[str, str, bytes]:
service = app.state.services
source_id = "local-source"
video_id = "local-video"
frame_id = "local-frame"
video_bytes = b"0123456789abcdef"
video_path = tmp_path / "sample.mp4"
video_path.write_bytes(video_bytes)
now = utcnow()
with service.db.transaction() as conn:
conn.execute(
"INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)",
(source_id, "local", "本地资料库", json.dumps({"path": str(tmp_path)}), now, now),
)
conn.execute(
"INSERT INTO videos(id,source_id,source_key,display_name,location,fingerprint,duration_ms,status,"
"available,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)",
(
video_id,
source_id,
"sample.mp4",
"海边假期.mp4",
str(video_path),
"fingerprint",
120_000,
"ready",
1,
now,
now,
),
)
conn.execute(
"INSERT INTO frames(id,video_id,timestamp_ms,segment_start_ms,segment_end_ms,thumbnail_path,created_at) "
"VALUES(?,?,?,?,?,?,?)",
(frame_id, video_id, 42_000, 40_000, 48_000, str(tmp_path / "frame.webp"), now),
)
tokens = " ".join(search_tokens("海边日落"))
for kind in ("subtitle", "filename"):
entry_id = f"text-{kind}"
conn.execute(
"INSERT INTO text_entries(id,video_id,frame_id,kind,start_ms,end_ms,raw_text,tokens,created_at) "
"VALUES(?,?,?,?,?,?,?,?,?)",
(entry_id, video_id, frame_id, kind, 40_000, 48_000, "海边日落", tokens, now),
)
conn.execute("INSERT INTO text_fts(entry_id,tokens) VALUES(?,?)", (entry_id, tokens))
audio_tokens = " ".join(search_tokens("海浪声音"))
conn.execute(
"INSERT INTO text_entries(id,video_id,frame_id,kind,start_ms,end_ms,raw_text,tokens,created_at) "
"VALUES('text-audio',?,?, 'audio',?,?,?,?,?)",
(video_id, frame_id, 65_000, 68_000, "远处传来海浪声音", audio_tokens, now),
)
conn.execute("INSERT INTO text_fts(entry_id,tokens) VALUES('text-audio',?)", (audio_tokens,))
_, token = service.auth.create_api_token("test")
return token, video_id, video_bytes
def test_gateway_media_token_reads_native_media_without_cookie_or_session_header(tmp_path: Path):
socket_path = tmp_path / "imagefind.sock"
settings = Settings(
data_dir=tmp_path / "data",
gateway_socket=socket_path,
gateway_prefix="/app/imagefind",
embedding_backend="hash",
scan_interval_seconds=86400,
)
settings.prepare()
app = create_app(settings)
service = app.state.services
service.auth.setup("gateway native media test password")
_, video_id, video_bytes = _seed_local_search(app, tmp_path)
thumbnail_path = tmp_path / "frame.webp"
thumbnail_path.write_bytes(b"webp-thumbnail")
preview = settings.preview_dir / "media123"
preview.mkdir(parents=True)
(preview / "index.m3u8").write_text("#EXTM3U\nsegment-00001.ts\n", encoding="utf-8")
(preview / "segment-00001.ts").write_bytes(b"gateway-segment")
class UnixSocketScope:
async def __call__(self, scope, receive, send):
mounted = dict(scope)
mounted["server"] = (str(socket_path), None)
await app(mounted, receive, send)
identity = {
"X-Trim-Isadmin": "true",
"X-Trim-Userid": "nas-admin",
"X-Trim-Username": "Administrator",
}
gateway_transport = httpx.ASGITransport(app=UnixSocketScope())
async def scenario():
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers=identity,
) as login_client:
login = await login_client.post("/app/imagefind/api/v1/auth/gateway")
assert login.status_code == 200
session_token = login.json()["gateway_session_token"]
media_token = login.json()["gateway_media_token"]
assert media_token and media_token != session_token
refreshed = await login_client.post(
"/app/imagefind/api/v1/auth/gateway/media-token",
headers={"X-CSRF-Token": login.json()["csrf_token"]},
)
assert refreshed.status_code == 200
refreshed_token = refreshed.json()["gateway_media_token"]
query = f"media_token={media_token}"
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers=identity,
) as native:
thumbnail = await native.get(f"/app/imagefind/api/v1/frames/local-frame/thumbnail?{query}")
assert thumbnail.status_code == 200
assert thumbnail.content == b"webp-thumbnail"
stream = await native.get(
f"/app/imagefind/api/v1/videos/{video_id}/stream?{query}",
headers={"Range": "bytes=2-5"},
)
assert stream.status_code == 206
assert stream.content == video_bytes[2:6]
download = await native.get(f"/app/imagefind/api/v1/videos/{video_id}/download?{query}")
assert download.status_code == 200
assert download.content == video_bytes
playlist = await native.get(f"/app/imagefind/api/v1/previews/media123/index.m3u8?{query}")
assert playlist.status_code == 200
assert f"segment-00001.ts?media_token={media_token}" in playlist.text
segment = await native.get(
f"/app/imagefind/api/v1/previews/media123/segment-00001.ts?{query}"
)
assert segment.status_code == 200
assert segment.content == b"gateway-segment"
assert (await native.get(f"/app/imagefind/api/v1/sources?{query}")).status_code == 401
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers={**identity, "X-Trim-Userid": "another-admin"},
) as wrong_user:
assert (
await wrong_user.get(f"/app/imagefind/api/v1/videos/{video_id}/stream?{query}")
).status_code == 401
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://test"
) as direct:
assert (await direct.get(f"/api/v1/videos/{video_id}/stream?{query}")).status_code == 401
with service.db.transaction() as conn:
conn.execute(
"UPDATE gateway_media_tokens SET expires_at='2000-01-01T00:00:00+00:00' WHERE token_hash=?",
(service.auth._digest(refreshed_token),),
)
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers=identity,
) as expired:
assert (
await expired.get(
f"/app/imagefind/api/v1/videos/{video_id}/stream?media_token={refreshed_token}"
)
).status_code == 401
service.auth.logout(session_token)
async with httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers=identity,
) as revoked:
assert (
await revoked.get(f"/app/imagefind/api/v1/videos/{video_id}/stream?{query}")
).status_code == 401
asyncio.run(scenario())
def test_search_match_sources_local_range_and_hls_paths(tmp_path: Path):
settings = Settings(data_dir=tmp_path, embedding_backend="hash", scan_interval_seconds=86400)
settings.prepare()
app = create_app(settings)
token, video_id, video_bytes = _seed_local_search(app, tmp_path)
headers = {"Authorization": f"Bearer {token}"}
preview = settings.preview_dir / "abc123"
preview.mkdir(parents=True)
(preview / "index.m3u8").write_bytes(b"#EXTM3U\n")
(preview / "segment-00001.ts").write_bytes(b"segment")
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
search = await client.post("/api/v1/search", headers=headers, json={"text": "海边日落"})
assert search.status_code == 200
item = search.json()["items"][0]
assert item["video_id"] == video_id
assert item["timestamp_ms"] == 40_000
assert item["match_sources"] == ["filename", "subtitle"]
assert {detail["type"] for detail in item["match_details"]} == {"metadata", "subtitle"}
audio = await client.post(
"/api/v1/search",
headers=headers,
json={"text": "海浪声音", "recognition_types": ["audio"]},
)
audio_item = audio.json()["items"][0]
assert audio_item["match_sources"] == ["audio"]
assert audio_item["segment_start_ms"] == 65_000
assert audio_item["match_details"][0]["text"] == "远处传来海浪声音"
excluded = await client.post(
"/api/v1/search",
headers=headers,
json={"text": "海浪声音", "recognition_types": ["ocr"]},
)
assert excluded.json()["items"] == []
partial = await client.get(
f"/api/v1/videos/{video_id}/stream",
headers={**headers, "Range": "bytes=2-5"},
)
assert partial.status_code == 206
assert partial.content == video_bytes[2:6]
assert partial.headers["content-range"] == f"bytes 2-5/{len(video_bytes)}"
assert partial.headers["accept-ranges"] == "bytes"
suffix = await client.get(
f"/api/v1/videos/{video_id}/stream",
headers={**headers, "Range": "bytes=-3"},
)
assert suffix.status_code == 206
assert suffix.content == video_bytes[-3:]
invalid = await client.get(
f"/api/v1/videos/{video_id}/stream",
headers={**headers, "Range": "bytes=99-100"},
)
assert invalid.status_code == 416
playlist = await client.get("/api/v1/previews/abc123/index.m3u8", headers=headers)
assert playlist.status_code == 200
assert playlist.content == b"#EXTM3U\n"
segment = await client.get("/api/v1/previews/abc123/segment-00001.ts", headers=headers)
assert segment.status_code == 200
assert segment.content == b"segment"
assert (await client.get("/api/v1/previews/not-safe/index.m3u8", headers=headers)).status_code == 404
assert (await client.get("/api/v1/previews/abc123/metadata.json", headers=headers)).status_code == 404
asyncio.run(scenario())
def test_audio_search_has_an_independent_candidate_budget_and_simplified_traditional_variants(
tmp_path: Path,
):
settings = Settings(data_dir=tmp_path, embedding_backend="hash", scan_interval_seconds=86400)
settings.prepare()
app = create_app(settings)
token, video_id, _ = _seed_local_search(app, tmp_path)
service = app.state.services
now = utcnow()
simplified = "繁体关键词"
traditional = "繁體關鍵詞"
with service.db.transaction() as conn:
for index in range(520):
tokens = " ".join(search_tokens(traditional))
entry_id = f"crowding-ocr-{index:03d}"
conn.execute(
"INSERT INTO text_entries(id,video_id,frame_id,kind,start_ms,end_ms,raw_text,tokens,created_at) "
"VALUES(?,?,?,'ocr',0,1000,?,?,?)",
(entry_id, video_id, "local-frame", traditional, tokens, now),
)
tokens = " ".join(search_tokens(simplified))
conn.execute(
"INSERT INTO text_entries(id,video_id,frame_id,kind,start_ms,end_ms,raw_text,tokens,created_at) "
"VALUES('audio-simplified',?,?, 'audio',66000,69000,?,?,?)",
(video_id, "local-frame", f"这里说的是{simplified}", tokens, now),
)
async def scenario():
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response = await client.post(
"/api/v1/search",
headers={"Authorization": f"Bearer {token}"},
json={"text": traditional},
)
assert response.status_code == 200
item = response.json()["items"][0]
assert "audio" in item["match_sources"]
assert any(detail["type"] == "audio" for detail in item["match_details"])
asyncio.run(scenario())
def test_transcript_coverage_pagination_and_manual_reindex(tmp_path: Path, monkeypatch):
settings = Settings(data_dir=tmp_path, embedding_backend="hash", scan_interval_seconds=86400)
settings.prepare()
app = create_app(settings)
token, video_id, _ = _seed_local_search(app, tmp_path)
service = app.state.services
monkeypatch.setattr(service.models, "runnable_component_versions", lambda: {"audio": "audio-v1"})
now = utcnow()
with service.db.transaction() as conn:
conn.execute(
"UPDATE videos SET audio_model_version='audio-v1',audio_index_revision=?,"
"audio_detected_language='zh',audio_quality_score=.94,audio_rejected_segments=1,"
"audio_quality_flags_json='not-json' WHERE id=?",
(SPEECH_INDEX_REVISION, video_id),
)
conn.execute("DELETE FROM text_entries WHERE video_id=? AND kind='audio'", (video_id,))
for index in range(35):
text = f"{index + 1}个中文音频片段"
conn.execute(
"INSERT INTO text_entries(id,video_id,frame_id,kind,start_ms,end_ms,raw_text,tokens,created_at) "
"VALUES(?,?,NULL,'audio',?,?,?,?,?)",
(
f"transcript-{index:02d}",
video_id,
index * 1000,
index * 1000 + 900,
text,
" ".join(search_tokens(text)),
now,
),
)
# A trusted legacy transcript remains searchable while the current model reindex is queued.
conn.execute("UPDATE videos SET audio_model_version='audio-legacy' WHERE id=?", (video_id,))
async def scenario():
headers = {"Authorization": f"Bearer {token}"}
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
coverage = await client.get("/api/v1/search/coverage", headers=headers)
assert coverage.status_code == 200
coverage_payload = coverage.json()
assert coverage_payload["indexed"] == 0
assert coverage_payload["searchable"] == 1
assert coverage_payload["empty"] == 0
assert coverage_payload["percent"] == 100.0
assert coverage_payload["current_model_percent"] == 0.0
with service.db.transaction() as conn:
conn.execute(
"UPDATE videos SET audio_model_version='audio-v1' WHERE id=?", (video_id,)
)
current_coverage = await client.get("/api/v1/search/coverage", headers=headers)
assert current_coverage.status_code == 200
assert current_coverage.json()["indexed"] == 1
assert current_coverage.json()["searchable"] == 1
assert current_coverage.json()["current_model_percent"] == 100.0
transcript = await client.get(
f"/api/v1/videos/{video_id}/transcript?page=2&page_size=10", headers=headers
)
assert transcript.status_code == 200
payload = transcript.json()
assert payload["status"] == "ready"
assert payload["page"] == 2 and payload["pages"] == 4 and payload["total"] == 35
assert payload["items"][0]["start_ms"] == 10_000
assert payload["detected_language"] == "zh"
assert payload["quality_state"] == "ready"
assert payload["quality_flags"] == []
speech = await client.get("/api/v1/speech/config", headers=headers)
assert speech.status_code == 200
assert speech.json()["language_policy"] == "zh_priority"
configured = await client.patch(
"/api/v1/speech/config",
headers=headers,
json={"language_policy": "auto", "quality_profile": "balanced"},
)
assert configured.status_code == 200
assert configured.json()["quality_profile"] == "balanced"
queued = await client.post(
f"/api/v1/videos/{video_id}/transcript/reindex", headers=headers, json={"language": "zh"}
)
assert queued.status_code == 202
with service.db.read() as conn:
job = conn.execute(
"SELECT kind,priority,payload_json FROM jobs WHERE id=?", (queued.json()["job_id"],)
).fetchone()
assert {"kind": job["kind"], "priority": job["priority"]} == {
"kind": "transcribe_audio",
"priority": 0,
}
assert json.loads(job["payload_json"]) == {"video_id": video_id, "language": "zh"}
asyncio.run(scenario())
def test_webdav_range_proxy_forwards_headers_and_hides_credentials(tmp_path: Path, monkeypatch):
socket_path = tmp_path / "imagefind.sock"
settings = Settings(
data_dir=tmp_path,
gateway_socket=socket_path,
gateway_prefix="/app/imagefind",
embedding_backend="hash",
scan_interval_seconds=86400,
)
settings.prepare()
app = create_app(settings)
service = app.state.services
service.auth.setup("gateway WebDAV media test password")
now = utcnow()
secret_blob = service.secrets.encrypt_json({"password": "remote-password"})
with service.db.transaction() as conn:
conn.execute(
"INSERT INTO sources(id,kind,name,config_json,secret_blob,created_at,updated_at) "
"VALUES(?,?,?,?,?,?,?)",
(
"dav-source",
"webdav",
"远程资料库",
json.dumps(
{
"base_url": "https://dav.example/videos/",
"username": "remote-user",
"verify_tls": True,
}
),
secret_blob,
now,
now,
),
)
conn.execute(
"INSERT INTO videos(id,source_id,source_key,display_name,location,fingerprint,status,available,"
"created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?)",
(
"dav-video",
"dav-source",
"movie.mp4",
"movie.mp4",
"https://dav.example/videos/movie.mp4",
"remote-fingerprint",
"ready",
1,
now,
now,
),
)
_, token = service.auth.create_api_token("test")
service.db.set_setting(
"webdav_server", {"enabled": True, "source_id": "dav-source", "relative_path": ""}
)
webdav_auth = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
observed: dict[str, object] = {"requests": []}
class UnixSocketScope:
async def __call__(self, scope, receive, send):
mounted = dict(scope)
mounted["server"] = (str(socket_path), None)
await app(mounted, receive, send)
identity = {
"X-Trim-Isadmin": "true",
"X-Trim-Userid": "nas-admin",
"X-Trim-Username": "Administrator",
}
class FakeAsyncClient:
def __init__(self, **kwargs):
observed["client_kwargs"] = kwargs
def build_request(self, method, url, headers):
request = httpx.Request(method, url, headers=headers)
observed["requests"].append(request)
return request
async def send(self, request, stream=False):
assert stream is True
if request.method == "HEAD":
return httpx.Response(405, request=request)
if request.headers.get("range") == "bytes=0-0":
return httpx.Response(
206,
request=request,
content=b"0",
headers={
"Content-Length": "1",
"Content-Range": "bytes 0-0/10",
"Accept-Ranges": "bytes",
"Content-Type": "application/octet-stream",
},
)
return httpx.Response(
206,
request=request,
content=b"2345",
headers={
"Content-Length": "4",
"Content-Range": "bytes 2-5/10",
"Accept-Ranges": "bytes",
"Content-Type": "application/octet-stream",
},
)
async def aclose(self):
observed["closed"] = True
async def scenario():
transport = httpx.ASGITransport(app=app)
gateway_transport = httpx.ASGITransport(app=UnixSocketScope())
async with (
httpx.AsyncClient(transport=transport, base_url="http://test") as client,
httpx.AsyncClient(
transport=gateway_transport,
base_url="http://test",
headers=identity,
) as gateway,
):
login = await gateway.post("/app/imagefind/api/v1/auth/gateway")
assert login.status_code == 200
media_token = login.json()["gateway_media_token"]
monkeypatch.setattr(api_module.httpx, "AsyncClient", FakeAsyncClient)
response = await client.get(
"/api/v1/videos/dav-video/stream",
headers={"Authorization": f"Bearer {token}", "Range": "bytes=2-5", "If-Range": '"etag"'},
)
assert response.status_code == 206
assert response.content == b"2345"
assert response.headers["content-range"] == "bytes 2-5/10"
assert response.headers["content-type"] == "video/mp4"
gateway_response = await gateway.get(
f"/app/imagefind/api/v1/videos/dav-video/stream?media_token={media_token}",
headers={"Range": "bytes=2-5", "If-Range": '"etag"'},
)
assert gateway_response.status_code == 206
assert gateway_response.content == b"2345"
assert gateway_response.headers["content-type"] == "video/mp4"
head = await client.request(
"HEAD",
"/webdav/movie.mp4",
headers=webdav_auth,
)
assert head.status_code == 200
assert head.headers["content-length"] == "10"
assert "content-range" not in head.headers
assert head.headers["accept-ranges"] == "bytes"
assert head.headers["content-type"] == "video/mp4"
ranged_head = await client.request(
"HEAD",
"/webdav/movie.mp4",
headers={**webdav_auth, "Range": "bytes=2-5"},
)
assert ranged_head.status_code == 206
assert ranged_head.headers["content-length"] == "4"
assert ranged_head.headers["content-range"] == "bytes 2-5/10"
asyncio.run(scenario())
requests = observed["requests"]
assert isinstance(requests, list)
assert [request.method for request in requests] == ["GET", "GET", "HEAD", "GET", "HEAD", "GET"]
upstream = requests[0]
assert isinstance(upstream, httpx.Request)
assert upstream.url == "https://dav.example/videos/movie.mp4"
assert upstream.headers["range"] == "bytes=2-5"
assert upstream.headers["if-range"] == '"etag"'
assert "remote-password" not in str(upstream.url)
assert observed["closed"] is True
File diff suppressed because it is too large Load Diff
+711
View File
@@ -0,0 +1,711 @@
import asyncio
import json
import shutil
import sqlite3
import stat
from pathlib import Path
import httpx
import pytest
from imagefind.backups import (
HEADER,
BackupError,
BackupLimitError,
BackupNotEmptyError,
BackupService,
BackupStorageError,
)
from imagefind.config import Settings
from imagefind.container import Services
from imagefind.database import utcnow
from imagefind.main import create_app
PASSWORD = "a separate backup password"
def make_services(path: Path, **overrides) -> Services:
settings = Settings(
data_dir=path,
embedding_backend="hash",
scan_interval_seconds=86400,
backup_reserve_mb=0,
**overrides,
)
settings.prepare()
return Services(settings)
def seed_source(services: Services, root: Path, *, source_id: str = "source-1", enabled: int = 1) -> None:
now = utcnow()
secret_blob = services.secrets.encrypt_json(
{"password": "remote login", "crypt_password": "rclone crypt", "crypt_salt": "rclone salt"}
)
with services.db.transaction() as conn:
conn.execute(
"INSERT INTO sources(id,kind,name,config_json,secret_blob,enabled,created_at,updated_at) "
"VALUES(?,?,?,?,?,?,?,?)",
(
source_id,
"webdav",
"Encrypted AList",
json.dumps(
{
"driver": "alist",
"mode": "encrypted",
"base_url": "https://alist.example",
"root_path": "private",
"username": "backup-user",
"verify_tls": True,
}
),
secret_blob,
enabled,
now,
now,
),
)
def seed_full(services: Services, media_root: Path) -> None:
now = utcnow()
media_root.mkdir(parents=True)
video_path = media_root / "movie.mp4"
video_path.write_bytes(b"video")
with services.db.transaction() as conn:
conn.execute(
"INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)",
("local-source", "local", "Local", json.dumps({"path": str(media_root)}), now, now),
)
conn.execute(
"INSERT INTO videos(id,source_id,source_key,display_name,location,size_bytes,fingerprint,duration_ms,"
"width,height,codec,container,status,available,indexed_fingerprint,basic_fingerprint,"
"visual_model_version,ocr_model_version,faces_model_version,error,created_at,updated_at) "
"VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
(
"video-1",
"local-source",
"movie.mp4",
"movie.mp4",
str(video_path),
5,
"old-fingerprint",
42_000,
1920,
1080,
"h264",
"mp4",
"ready",
1,
"indexed",
"basic",
"visual-v1",
"ocr-v1",
"faces-v1",
"old error",
now,
now,
),
)
conn.execute(
"INSERT INTO people(id,name,normalized_name,is_named,created_at,updated_at) VALUES(?,?,?,?,?,?)",
("person-1", "Named Person", "named person", 1, now, now),
)
conn.execute(
"INSERT INTO actors(id,name,aliases_json,person_id,created_at,updated_at) VALUES(?,?,?,?,?,?)",
("actor-1", "Actor", '["Alias"]', "person-1", now, now),
)
conn.execute(
"INSERT INTO tag_groups(id,name,selection_mode,sort_order,created_at,updated_at) VALUES(?,?,?,?,?,?)",
("group-1", "Genre", "multi", 1, now, now),
)
conn.execute(
"INSERT INTO tags(id,group_id,name,ai_enabled,ai_method,ai_description,ai_threshold,match_terms_json,"
"created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?)",
("tag-1", "group-1", "Favorite", 1, "text", "", 0.5, '["movie"]', now, now),
)
conn.execute(
"INSERT INTO video_metadata(video_id,title,catalog_code,studio,series,release_date,description,updated_at) "
"VALUES(?,?,?,?,?,?,?,?)",
("video-1", "Manual title", "CAT-001", "Studio", "Series", "2026-01-01", "Notes", now),
)
conn.execute("INSERT INTO video_actors(video_id,actor_id) VALUES('video-1','actor-1')")
conn.execute("INSERT INTO video_tags(video_id,tag_id) VALUES('video-1','tag-1')")
conn.execute(
"INSERT INTO video_state(video_id,liked,favorited,progress_ms,completed,last_played_at,updated_at) "
"VALUES('video-1',1,1,12345,0,?,?)",
(now, now),
)
conn.execute(
"INSERT INTO frames(id,video_id,timestamp_ms,segment_start_ms,segment_end_ms,thumbnail_path,vector_blob,"
"created_at) VALUES('frame-1','video-1',1000,0,2000,'thumb.webp',?,?)",
(b"vector", now),
)
conn.execute(
"INSERT INTO faces(id,frame_id,video_id,person_id,bbox_json,confidence,vector_blob,created_at) "
"VALUES('face-1','frame-1','video-1','person-1','[]',0.9,?,?)",
(b"face-vector", now),
)
conn.execute(
"INSERT INTO text_entries(id,video_id,frame_id,kind,raw_text,tokens,created_at) "
"VALUES('text-1','video-1','frame-1','ocr','secret text','secret text',?)",
(now,),
)
conn.execute("INSERT INTO text_fts(entry_id,tokens) VALUES('text-1','secret text')")
conn.execute(
"INSERT INTO tag_suggestions(id,video_id,tag_id,confidence,video_fingerprint,tag_revision,created_at,"
"updated_at) VALUES('suggestion-1','video-1','tag-1',0.9,'old','old',?,?)",
(now, now),
)
services.db.set_setting(
"preferences",
{
"autoplay": False,
"mask_covers": True,
"home_video_columns": 1,
"upload_paths": {"local-source": "incoming"},
},
)
services.db.set_setting("model_hf_endpoint", "https://hf-mirror.example")
services.db.set_setting("pip_index_url", "https://pypi-mirror.example/simple")
services.db.set_setting("pytorch_index_url", "https://torch-mirror.example/cpu")
services.db.set_setting("profile_nickname", "Backup Administrator")
def test_keys_backup_restores_sources_with_a_new_secret_key(tmp_path: Path):
source = make_services(tmp_path / "old")
source.auth.setup("old administrator password")
seed_source(source, tmp_path)
source.auth.create_api_token("old token")
artifact = source.backups.export("keys", PASSWORD)
assert stat.S_IMODE(artifact.path.stat().st_mode) == 0o600
assert list(artifact.temporary_dir.iterdir()) == [artifact.path]
target = make_services(tmp_path / "new")
target.auth.setup("new administrator password")
target.auth.create_api_token("new token")
old_secret_key = (source.settings.data_dir / ".secret-key").read_bytes()
new_secret_key = (target.settings.data_dir / ".secret-key").read_bytes()
assert old_secret_key != new_secret_key
result = target.backups.restore(artifact.path, PASSWORD, True)
restored = target.sources.get("source-1")
assert restored["secrets"] == {
"password": "remote login",
"crypt_password": "rclone crypt",
"crypt_salt": "rclone salt",
}
assert result["scope"] == "keys"
assert len(result["scan_jobs"]) == 1
assert target.auth.login("new administrator password")
with target.db.read() as conn:
assert conn.execute("SELECT count(*) FROM admin").fetchone()[0] == 1
assert conn.execute("SELECT count(*) FROM api_tokens").fetchone()[0] == 1
assert conn.execute("SELECT count(*) FROM videos").fetchone()[0] == 0
artifact.cleanup()
def test_full_backup_keeps_user_metadata_and_excludes_derived_data(tmp_path: Path):
media_root = tmp_path / "media"
source = make_services(tmp_path / "old")
source.auth.setup("old administrator password")
seed_full(source, media_root)
source.auth.create_api_token("old token")
artifact = source.backups.export("full", PASSWORD)
target = make_services(tmp_path / "new")
target.auth.setup("new administrator password")
target.auth.create_api_token("new token")
result = target.backups.restore(artifact.path, PASSWORD, True)
assert result["scope"] == "full"
assert result["counts"]["videos"] == 1
assert result["counts"]["actors"] == 1
with target.db.read() as conn:
video = conn.execute("SELECT * FROM videos WHERE id='video-1'").fetchone()
assert dict(video) | {} # sqlite.Row remains readable after the assertion block
assert video["source_key"] == "movie.mp4"
assert video["location"] == ""
assert video["size_bytes"] == 0
assert video["status"] == "pending"
assert video["available"] == 0
assert video["indexed_fingerprint"] is None
assert video["basic_fingerprint"] is None
assert video["visual_model_version"] is None
assert video["ocr_model_version"] is None
assert video["faces_model_version"] is None
metadata = conn.execute("SELECT * FROM video_metadata WHERE video_id='video-1'").fetchone()
assert metadata["title"] == "Manual title"
state = conn.execute("SELECT * FROM video_state WHERE video_id='video-1'").fetchone()
assert (state["liked"], state["favorited"], state["progress_ms"]) == (1, 1, 12345)
actor = conn.execute("SELECT * FROM actors WHERE id='actor-1'").fetchone()
assert actor["person_id"] is None
assert conn.execute("SELECT count(*) FROM video_actors").fetchone()[0] == 1
assert conn.execute("SELECT count(*) FROM video_tags").fetchone()[0] == 1
for table in ("frames", "text_entries", "people", "faces", "tag_suggestions"):
assert conn.execute(f"SELECT count(*) FROM {table}").fetchone()[0] == 0
assert conn.execute("SELECT count(*) FROM api_tokens").fetchone()[0] == 1
assert target.db.setting("preferences")["mask_covers"] is True
assert target.db.setting("preferences")["home_video_columns"] == 1
assert target.db.setting("model_hf_endpoint") == "https://hf-mirror.example"
assert target.db.setting("pip_index_url") == "https://pypi-mirror.example/simple"
assert target.db.setting("pytorch_index_url") == "https://torch-mirror.example/cpu"
assert target.db.setting("profile_nickname") == "Backup Administrator"
assert target.settings.model_hf_endpoint == "https://hf-mirror.example"
assert target.settings.pip_index_url == "https://pypi-mirror.example/simple"
assert target.settings.pytorch_index_url == "https://torch-mirror.example/cpu"
artifact.cleanup()
def test_full_backup_restores_collection_membership_cover_and_order(tmp_path: Path):
source = make_services(tmp_path / "old")
seed_full(source, tmp_path / "media")
now = utcnow()
with source.db.transaction() as conn:
conn.execute(
"INSERT INTO collections(id,name,description,cover_video_id,created_at,updated_at) "
"VALUES(?,?,?,?,?,?)",
("collection-1", "Weekend", "Saved for the weekend", "video-1", now, now),
)
conn.execute(
"INSERT INTO collection_videos(collection_id,video_id,position,added_at) VALUES(?,?,?,?)",
("collection-1", "video-1", 7, now),
)
conn.execute(
"INSERT INTO collection_items(id,collection_id,parent_id,kind,name,video_id,position,"
"created_at,updated_at) "
"VALUES('chapter-1','collection-1',NULL,'group','第一章',NULL,0,?,?)",
(now, now),
)
conn.execute(
"INSERT INTO collection_items(id,collection_id,parent_id,kind,name,video_id,position,"
"created_at,updated_at) "
"VALUES('chapter-video','collection-1','chapter-1','video',NULL,'video-1',0,?,?)",
(now, now),
)
conn.execute(
"INSERT INTO collection_tags(collection_id,tag_id,created_at) VALUES(?,?,?)",
("collection-1", "tag-1", now),
)
conn.execute(
"INSERT INTO video_markers(id,video_id,position_ms,title,created_at,updated_at) "
"VALUES(?,?,?,?,?,?)",
("marker-1", "video-1", 12_345, "精彩片段", now, now),
)
conn.execute(
"INSERT INTO video_tombstones(source_id,source_key,video_id,display_name,source_deleted,deleted_at) "
"VALUES(?,?,?,?,?,?)",
("local-source", "removed.mp4", "removed-video", "removed.mp4", 0, now),
)
source.db.set_setting(
"webdav_server",
{"enabled": True, "source_id": "local-source", "relative_path": "WebDAV 上传"},
)
artifact = source.backups.export("full", PASSWORD)
target = make_services(tmp_path / "new")
result = target.backups.restore(artifact.path, PASSWORD, True)
assert result["counts"]["collections"] == 1
assert result["counts"]["collection_videos"] == 1
assert result["counts"]["collection_items"] == 2
assert result["counts"]["collection_tags"] == 1
assert result["counts"]["video_markers"] == 1
assert result["counts"]["video_tombstones"] == 1
with target.db.read() as conn:
collection = conn.execute("SELECT * FROM collections WHERE id='collection-1'").fetchone()
membership = conn.execute(
"SELECT * FROM collection_videos WHERE collection_id='collection-1' AND video_id='video-1'"
).fetchone()
assert collection["name"] == "Weekend"
assert collection["description"] == "Saved for the weekend"
assert collection["cover_video_id"] == "video-1"
assert membership["position"] == 7
restored_item = conn.execute(
"SELECT parent_id,kind,video_id FROM collection_items WHERE id='chapter-video'"
).fetchone()
assert dict(restored_item) == {
"parent_id": "chapter-1",
"kind": "video",
"video_id": "video-1",
}
assert conn.execute(
"SELECT tag_id FROM collection_tags WHERE collection_id='collection-1'"
).fetchone()[0] == "tag-1"
marker = conn.execute("SELECT * FROM video_markers WHERE id='marker-1'").fetchone()
assert marker["position_ms"] == 12_345
assert marker["title"] == "精彩片段"
tombstone = conn.execute(
"SELECT * FROM video_tombstones WHERE source_id='local-source' AND source_key='removed.mp4'"
).fetchone()
assert tombstone["video_id"] == "removed-video"
assert tombstone["source_deleted"] == 0
assert target.db.setting("webdav_server") == {
"enabled": True,
"source_id": "local-source",
"relative_path": "WebDAV 上传",
}
artifact.cleanup()
@pytest.mark.parametrize("mutation", ["password", "tampered", "truncated", "version"])
def test_invalid_encrypted_backup_never_changes_target(tmp_path: Path, mutation: str):
source = make_services(tmp_path / f"old-{mutation}")
seed_source(source, tmp_path)
artifact = source.backups.export("keys", PASSWORD)
candidate = tmp_path / f"{mutation}.ifbackup"
data = bytearray(artifact.path.read_bytes())
password = PASSWORD
if mutation == "password":
password = "the incorrect backup password"
elif mutation == "tampered":
data[HEADER.size + 1] ^= 1
elif mutation == "truncated":
data = data[:-20]
else:
data[8:10] = (999).to_bytes(2, "big")
candidate.write_bytes(data)
target = make_services(tmp_path / f"new-{mutation}")
target.auth.setup("new administrator password")
with pytest.raises(BackupError):
target.backups.restore(candidate, password, True)
assert target.backups.status()["empty"] is True
assert target.auth.login("new administrator password")
artifact.cleanup()
def corrupt_logical_backup(
service: BackupService,
artifact_path: Path,
destination: Path,
statement: str,
) -> None:
work = destination.parent / f"work-{destination.stem}"
work.mkdir()
compressed = work / "payload.gz"
logical = work / "logical.sqlite3"
recompressed = work / "changed.gz"
service._decrypt(artifact_path, compressed, PASSWORD)
service._decompress(compressed, logical)
connection = sqlite3.connect(logical)
connection.execute(statement)
connection.commit()
connection.close()
service._compress(logical, recompressed)
service._encrypt(recompressed, destination, PASSWORD)
shutil.rmtree(work)
def downgrade_to_v1_logical_backup(
service: BackupService,
artifact_path: Path,
destination: Path,
) -> None:
work = destination.parent / f"work-{destination.stem}"
work.mkdir()
compressed = work / "payload.gz"
logical = work / "logical.sqlite3"
recompressed = work / "legacy.gz"
service._decrypt(artifact_path, compressed, PASSWORD)
service._decompress(compressed, logical)
connection = sqlite3.connect(logical)
counts = json.loads(connection.execute("SELECT counts_json FROM manifest WHERE id=1").fetchone()[0])
counts.pop("collections")
counts.pop("collection_videos")
counts.pop("collection_items")
counts.pop("collection_tags")
counts.pop("video_markers")
counts.pop("video_tombstones")
connection.execute("DROP TABLE video_tombstones")
connection.execute("DROP TABLE collection_tags")
connection.execute("DROP TABLE video_markers")
connection.execute("DROP TABLE collection_items")
connection.execute("DROP TABLE collection_videos")
connection.execute("DROP TABLE collections")
connection.execute(
"CREATE TABLE videos_v1(id TEXT PRIMARY KEY,source_id TEXT NOT NULL REFERENCES sources(id) "
"ON DELETE CASCADE,source_key TEXT NOT NULL,display_name TEXT NOT NULL,created_at TEXT NOT NULL,"
"updated_at TEXT NOT NULL,UNIQUE(source_id,source_key))"
)
connection.execute(
"INSERT INTO videos_v1 SELECT id,source_id,source_key,display_name,created_at,updated_at FROM videos"
)
connection.execute("DROP TABLE videos")
connection.execute("ALTER TABLE videos_v1 RENAME TO videos")
connection.execute(
"UPDATE manifest SET logical_schema_version=1,counts_json=? WHERE id=1",
(json.dumps(counts, sort_keys=True),),
)
connection.commit()
connection.close()
service._compress(logical, recompressed)
service._encrypt(recompressed, destination, PASSWORD)
shutil.rmtree(work)
def test_v1_full_backup_migrates_legacy_series_to_collection(tmp_path: Path):
source = make_services(tmp_path / "old")
seed_full(source, tmp_path / "media")
artifact = source.backups.export("full", PASSWORD)
legacy = tmp_path / "legacy-v1.ifbackup"
downgrade_to_v1_logical_backup(source.backups, artifact.path, legacy)
target = make_services(tmp_path / "new")
result = target.backups.restore(legacy, PASSWORD, True)
assert result["counts"].get("collections") is None
with target.db.read() as conn:
collection = conn.execute("SELECT id,name FROM collections").fetchone()
membership = conn.execute("SELECT collection_id,video_id,position FROM collection_videos").fetchone()
assert collection["name"] == "Series"
assert membership["collection_id"] == collection["id"]
assert membership["video_id"] == "video-1"
assert membership["position"] == 0
artifact.cleanup()
@pytest.mark.parametrize(
"statement",
[
"UPDATE videos SET source_id='missing-source'",
"UPDATE manifest SET counts_json='{}'",
],
)
def test_logical_relationship_or_count_corruption_is_rejected_atomically(tmp_path: Path, statement: str):
source = make_services(tmp_path / "old")
seed_full(source, tmp_path / "media")
artifact = source.backups.export("full", PASSWORD)
corrupted = tmp_path / f"corrupted-{abs(hash(statement))}.ifbackup"
corrupt_logical_backup(source.backups, artifact.path, corrupted, statement)
target = make_services(tmp_path / "new")
target.auth.setup("new administrator password")
with pytest.raises(BackupError):
target.backups.restore(corrupted, PASSWORD, True)
assert target.backups.status()["empty"] is True
artifact.cleanup()
def test_nonempty_system_limits_storage_warning_and_scan_relink(tmp_path: Path, monkeypatch):
media_root = tmp_path / "media"
source = make_services(tmp_path / "old")
seed_full(source, media_root)
artifact = source.backups.export("full", PASSWORD)
nonempty = make_services(tmp_path / "nonempty")
seed_source(nonempty, tmp_path, source_id="existing")
with pytest.raises(BackupNotEmptyError):
nonempty.backups.restore(artifact.path, PASSWORD, True)
limited = make_services(tmp_path / "limited", backup_upload_gb=0.0000001)
with pytest.raises(BackupLimitError):
limited.backups.restore(artifact.path, PASSWORD, True)
extract_limited = make_services(tmp_path / "extract-limited", backup_extract_gb=0.000001)
with pytest.raises(BackupLimitError):
extract_limited.backups.restore(artifact.path, PASSWORD, True)
target = make_services(tmp_path / "new")
result = target.backups.restore(artifact.path, PASSWORD, True)
assert result["warnings"] == []
target.scanner.scan(result["scan_jobs"][0]["job_id"], "local-source")
with target.db.read() as conn:
video = conn.execute("SELECT id,available,location FROM videos WHERE source_key='movie.mp4'").fetchone()
assert video["id"] == "video-1"
assert video["available"] == 1
assert video["location"] == str(media_root / "movie.mp4")
monkeypatch.setattr(shutil, "disk_usage", lambda _path: shutil._ntuple_diskusage(100, 100, 0))
with pytest.raises(BackupStorageError):
source.backups.ensure_free_space(1)
artifact.cleanup()
def test_restore_database_transaction_rolls_back_after_partial_insert(tmp_path: Path, monkeypatch):
source = make_services(tmp_path / "old")
seed_full(source, tmp_path / "media")
artifact = source.backups.export("full", PASSWORD)
target = make_services(tmp_path / "new")
target.auth.setup("current administrator password")
target.auth.create_api_token("current token")
original = target.backups._insert_rows
def fail_during_metadata(source_conn, destination_conn, table):
if table == "actors":
raise sqlite3.IntegrityError("injected transaction failure")
return original(source_conn, destination_conn, table)
monkeypatch.setattr(target.backups, "_insert_rows", fail_during_metadata)
with pytest.raises(sqlite3.IntegrityError, match="injected"):
target.backups.restore(artifact.path, PASSWORD, True)
assert target.backups.status()["empty"] is True
assert target.auth.login("current administrator password")
with target.db.read() as conn:
assert conn.execute("SELECT count(*) FROM sources").fetchone()[0] == 0
assert conn.execute("SELECT count(*) FROM tags").fetchone()[0] == 0
assert conn.execute("SELECT count(*) FROM api_tokens").fetchone()[0] == 1
artifact.cleanup()
def test_unavailable_local_source_restores_with_warning(tmp_path: Path):
source = make_services(tmp_path / "old")
now = utcnow()
with source.db.transaction() as conn:
conn.execute(
"INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)",
("missing-local", "local", "Missing disk", json.dumps({"path": str(tmp_path / "missing")}), now, now),
)
artifact = source.backups.export("keys", PASSWORD)
target = make_services(tmp_path / "new")
result = target.backups.restore(artifact.path, PASSWORD, True)
assert "路径暂不可用" in result["warnings"][0]
artifact.cleanup()
def test_backup_api_download_restore_and_nonempty_conflict(tmp_path: Path):
source_settings = Settings(
data_dir=tmp_path / "source-api",
embedding_backend="hash",
scan_interval_seconds=86400,
backup_reserve_mb=0,
)
source_settings.prepare()
source_app = create_app(source_settings)
seed_source(source_app.state.services, tmp_path)
async def export_scenario() -> bytes:
transport = httpx.ASGITransport(app=source_app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
setup = await client.post("/api/v1/setup", json={"password": "source administrator"})
csrf = setup.json()["csrf_token"]
status = await client.get("/api/v1/backups/status")
assert status.status_code == 200
assert status.json()["empty"] is False
response = await client.post(
"/api/v1/backups/export",
headers={"X-CSRF-Token": csrf},
json={"scope": "keys", "password": PASSWORD},
)
assert response.status_code == 200
assert response.headers["content-type"].startswith("application/vnd.imagefind.backup")
assert "imagefind-backup-" in response.headers["content-disposition"]
assert response.headers["x-imagefind-backup-scope"] == "keys"
return response.content
payload = asyncio.run(export_scenario())
target_settings = Settings(
data_dir=tmp_path / "target-api",
embedding_backend="hash",
scan_interval_seconds=86400,
backup_reserve_mb=0,
)
target_settings.prepare()
target_app = create_app(target_settings)
async def restore_scenario() -> None:
transport = httpx.ASGITransport(app=target_app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
setup = await client.post("/api/v1/setup", json={"password": "target administrator"})
csrf = setup.json()["csrf_token"]
status = await client.get("/api/v1/backups/status")
assert status.json()["can_restore"] is True
restored = await client.post(
"/api/v1/backups/restore",
headers={"X-CSRF-Token": csrf},
files={"file": ("portable.ifbackup", payload, "application/octet-stream")},
data={"password": PASSWORD, "confirmed": "true"},
)
assert restored.status_code == 200
assert restored.json()["counts"]["sources"] == 1
again = await client.post(
"/api/v1/backups/restore",
headers={"X-CSRF-Token": csrf},
files={"file": ("portable.ifbackup", payload, "application/octet-stream")},
data={"password": PASSWORD, "confirmed": "true"},
)
assert again.status_code == 409
asyncio.run(restore_scenario())
def test_async_backup_api_tracks_job_resource_download_and_clears_secret(tmp_path: Path, monkeypatch):
settings = Settings(
data_dir=tmp_path / "async-backup-api",
embedding_backend="hash",
scan_interval_seconds=86400,
backup_reserve_mb=0,
)
settings.prepare()
app = create_app(settings)
seed_source(app.state.services, tmp_path)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
setup = await client.post("/api/v1/setup", json={"password": "source administrator"})
csrf = setup.json()["csrf_token"]
queued = await client.post(
"/api/v1/backups",
headers={"X-CSRF-Token": csrf},
json={"scope": "keys", "password": PASSWORD},
)
assert queued.status_code == 202
export_id = queued.json()["id"]
job_id = queued.json()["job_id"]
with app.state.services.db.read() as conn:
row = conn.execute(
"SELECT status,secret_blob FROM backup_exports WHERE id=?", (export_id,)
).fetchone()
resource = conn.execute(
"SELECT resource_type,resource_id FROM job_resources WHERE job_id=?", (job_id,)
).fetchone()
assert row["status"] == "queued"
assert row["secret_blob"] and PASSWORD not in row["secret_blob"]
assert tuple(resource) == ("backup", export_id)
app.state.services.backups.run_export(job_id, export_id)
records = await client.get("/api/v1/backups")
completed = next(item for item in records.json() if item["id"] == export_id)
assert completed["status"] == "completed"
assert completed["size_bytes"] > 0
assert "secret_blob" not in completed
download = await client.get(f"/api/v1/backups/{export_id}/download")
assert download.status_code == 200
assert download.content.startswith(b"IFBACKUP")
assert download.headers["content-length"] == str(len(download.content))
failed = await client.post(
"/api/v1/backups",
headers={"X-CSRF-Token": csrf},
json={"scope": "keys", "password": PASSWORD},
)
failed_id = failed.json()["id"]
failed_job = failed.json()["job_id"]
def fail_export(*_args, **_kwargs):
raise BackupError("injected asynchronous export failure")
monkeypatch.setattr(app.state.services.backups, "export", fail_export)
with pytest.raises(BackupError, match="injected"):
app.state.services.backups.run_export(failed_job, failed_id)
with app.state.services.db.read() as conn:
failed_row = conn.execute(
"SELECT status,secret_blob,error FROM backup_exports WHERE id=?", (failed_id,)
).fetchone()
assert failed_row["status"] == "failed"
assert failed_row["secret_blob"] is None
assert "injected" in failed_row["error"]
unavailable = await client.get(f"/api/v1/backups/{failed_id}/download")
assert unavailable.status_code == 409
cancelled = app.state.services.backups.queue_export("keys", PASSWORD)
app.state.services.jobs.request_cancel(cancelled["job_id"])
with app.state.services.db.read() as conn:
cancelled_row = conn.execute(
"SELECT status,secret_blob,error FROM backup_exports WHERE id=?", (cancelled["id"],)
).fetchone()
assert cancelled_row["status"] == "failed"
assert cancelled_row["secret_blob"] is None
assert cancelled_row["error"] == "备份任务已取消"
asyncio.run(scenario())
+388
View File
@@ -0,0 +1,388 @@
import contextlib
import threading
import time
from pathlib import Path
import pytest
from imagefind.database import SCHEMA_VERSION, Database, DatabaseTransientError
from imagefind.security import AuthService, LoginRateLimitError, SecretStore
def test_database_instances_use_independent_postgres_transactions(tmp_path: Path):
path = tmp_path / "coordinated-postgres"
first = Database(path)
first.initialize()
second = Database(path)
writer_started = threading.Event()
release_writer = threading.Event()
second_finished = threading.Event()
errors: list[Exception] = []
def hold_writer():
try:
with first.transaction() as conn:
conn.execute("INSERT INTO settings(key,value,updated_at) VALUES('first','1','now')")
writer_started.set()
release_writer.wait(2)
except Exception as exc: # pragma: no cover - asserted below
errors.append(exc)
def queued_writer():
try:
writer_started.wait(2)
with second.transaction() as conn:
conn.execute("INSERT INTO settings(key,value,updated_at) VALUES('second','2','now')")
second_finished.set()
except Exception as exc: # pragma: no cover - asserted below
errors.append(exc)
holding = threading.Thread(target=hold_writer)
waiting = threading.Thread(target=queued_writer)
holding.start()
assert writer_started.wait(2)
waiting.start()
assert second_finished.wait(1)
assert first.status()["writer_queue_depth"] == 0
time.sleep(0.26)
release_writer.set()
holding.join(2)
waiting.join(2)
assert not errors
with first.read() as conn:
keys = {row[0] for row in conn.execute("SELECT key FROM settings WHERE key IN ('first','second')").fetchall()}
assert keys == {"first", "second"}
status = second.status()
assert status["engine"] == "postgresql"
assert status["journal_mode"] == "server"
assert status["writer_queue_depth"] == 0
assert status["writer_active"] is False
assert status["pool_max"] >= 2
def test_write_retry_recovers_a_transient_postgres_conflict(tmp_path: Path, monkeypatch):
db = Database(tmp_path / "retry-postgres")
db.initialize()
original_transaction = db.transaction
attempts = 0
@contextlib.contextmanager
def transient_transaction():
nonlocal attempts
attempts += 1
if attempts == 1:
raise DatabaseTransientError("transient PostgreSQL serialization failure")
with original_transaction() as conn:
yield conn
monkeypatch.setattr(db, "transaction", transient_transaction)
db.write_with_retry(
lambda conn: conn.execute("INSERT INTO settings(key,value,updated_at) VALUES('recovered','1','now')"),
timeout_seconds=2,
)
assert db.setting("recovered") == 1
assert attempts == 2
assert db.status()["lock_retry_count"] >= 1
def test_database_auth_and_encrypted_secrets(tmp_path: Path):
db = Database(tmp_path / "auth-postgres")
db.initialize()
auth = AuthService(db, session_days=1)
assert not auth.is_configured()
with pytest.raises(ValueError):
auth.setup("short")
auth.setup("correct horse battery staple")
token, csrf, _ = auth.login("correct horse battery staple")
assert auth.session(token)["csrf_token"] == csrf
assert auth.session(token)["auth_source"] == "local"
gateway_token, gateway_csrf, _ = auth.login_gateway("nas-admin", "Administrator")
gateway_session = auth.session(gateway_token)
assert gateway_session["csrf_token"] == gateway_csrf
assert gateway_session["auth_source"] == "gateway"
assert gateway_session["external_user_id"] == "nas-admin"
with pytest.raises(ValueError):
auth.login("wrong")
token_id, api_token = auth.create_api_token("automation")
assert auth.verify_api_token(api_token)
assert auth.verify_api_token(api_token)
with db.read() as conn:
assert conn.execute("SELECT last_used_at FROM api_tokens WHERE id=?", (token_id,)).fetchone()[0] is None
assert auth.flush_api_token_usage() == 1
with db.read() as conn:
assert conn.execute("SELECT last_used_at FROM api_tokens WHERE id=?", (token_id,)).fetchone()[0]
auth.set_password("a replacement password", replace=True)
assert auth.login("a replacement password")
assert auth.session(token) is None
assert auth.verify_api_token(api_token)
with pytest.raises(ValueError):
auth.login("correct horse battery staple")
auth.revoke_api_token(token_id)
assert not auth.verify_api_token(api_token)
store = SecretStore(tmp_path / "secret.key")
ciphertext = store.encrypt_json({"password": "not-plaintext"})
assert "not-plaintext" not in ciphertext
assert store.decrypt_json(ciphertext) == {"password": "not-plaintext"}
def test_login_rate_limit_is_scoped_by_client_and_success_clears_failures(tmp_path: Path, monkeypatch):
db = Database(tmp_path / "rate-limit-postgres")
db.initialize()
auth = AuthService(db, session_days=1)
auth.setup("correct horse battery staple")
monkeypatch.setattr(auth, "LOGIN_MAX_FAILURES", 2)
for _ in range(2):
with pytest.raises(ValueError, match="密码错误"):
auth.login("wrong password", client_key="192.0.2.10")
with pytest.raises(LoginRateLimitError) as blocked:
auth.login("correct horse battery staple", client_key="192.0.2.10")
assert blocked.value.retry_after > 0
# A separate client remains usable, and a successful login clears that
# client's partial failure history.
with pytest.raises(ValueError, match="密码错误"):
auth.login("wrong password", client_key="192.0.2.11")
token, _, _ = auth.login("correct horse battery staple", client_key="192.0.2.11")
assert auth.session(token)
with pytest.raises(ValueError, match="密码错误"):
auth.login("wrong password", client_key="192.0.2.11")
def test_api_token_scopes_are_enforced_and_legacy_admin_default_is_preserved(tmp_path: Path):
db = Database(tmp_path / "token-scopes-postgres")
db.initialize()
auth = AuthService(db, session_days=1)
_, webdav_token = auth.create_api_token("DAV only", ["webdav"])
_, media_token = auth.create_api_token("Media only", ["media:read"])
_, admin_token = auth.create_api_token("Administrator")
assert auth.verify_api_token(webdav_token, "webdav")
assert not auth.verify_api_token(webdav_token, "admin")
assert not auth.verify_api_token(webdav_token, "media:read")
assert auth.verify_api_token(media_token, "media:read")
assert not auth.verify_api_token(media_token, "webdav")
assert auth.verify_api_token(admin_token, "admin")
assert auth.verify_api_token(admin_token, "webdav")
assert auth.verify_api_token(admin_token, "media:read")
with pytest.raises(ValueError, match="权限范围"):
auth.create_api_token("invalid", ["unknown"])
def test_postgres_pool_is_bounded_and_reaps_extra_idle_connections(tmp_path: Path):
db = Database(tmp_path / "bounded-pool-postgres")
db._pool_min = 1
db._pool_max = 2
db._pool_timeout = 0.1
db._pool_idle_timeout = 1
db.initialize()
with db.read() as first, db.read() as second:
assert first.execute("SELECT 1").fetchone()[0] == 1
assert second.execute("SELECT 1").fetchone()[0] == 1
with pytest.raises(TimeoutError, match="连接超时"):
with db.read():
pass
status = db.status()
assert status["pool_size"] == 2
assert status["pool_wait_count"] >= 1
time.sleep(1.05)
with db.read() as conn:
assert conn.execute("SELECT 1").fetchone()[0] == 1
assert db.status()["pool_size"] == 1
activity = db.activity()
assert isinstance(activity["states"], dict)
assert activity["waiting"] >= 0
def test_password_validation_and_replacement_guard(tmp_path: Path):
db = Database(tmp_path / "password-postgres")
db.initialize()
auth = AuthService(db, session_days=1)
with pytest.raises(ValueError, match="至少"):
auth.set_password("short")
with pytest.raises(ValueError, match="超过"):
auth.set_password("x" * 257)
auth.set_password("initial administrator password")
with pytest.raises(ValueError, match="已经初始化"):
auth.set_password("second administrator password")
def test_postgres_mvcc_reader_stays_responsive_while_writer_is_open(tmp_path: Path):
db = Database(tmp_path / "mvcc-postgres")
db.initialize()
writer = db.connect()
try:
writer.execute("INSERT INTO settings(key,value,updated_at) VALUES('writer','1','now')")
started = time.monotonic()
with db.read() as reader:
assert reader.execute("SELECT count(*) FROM settings WHERE key='writer'").fetchone()[0] == 0
assert time.monotonic() - started < 0.5
finally:
writer.rollback()
writer.close()
def test_slow_reader_does_not_serialize_other_thread_readers(tmp_path: Path):
db = Database(tmp_path / "parallel-postgres")
db.initialize()
first_entered = threading.Event()
release_first = threading.Event()
second_finished = threading.Event()
errors: list[Exception] = []
def hold_reader():
try:
with db.read() as reader:
assert reader.execute("SELECT count(*) FROM settings").fetchone()[0] >= 1
first_entered.set()
assert release_first.wait(2)
except Exception as exc: # pragma: no cover - surfaced below
errors.append(exc)
def use_second_reader():
try:
assert first_entered.wait(2)
with db.read() as reader:
assert reader.execute("SELECT count(*) FROM settings").fetchone()[0] >= 1
second_finished.set()
except Exception as exc: # pragma: no cover - surfaced below
errors.append(exc)
first = threading.Thread(target=hold_reader)
second = threading.Thread(target=use_second_reader)
first.start()
second.start()
assert second_finished.wait(1), "an unrelated reader waited behind the first reader"
release_first.set()
first.join(timeout=2)
second.join(timeout=2)
assert errors == []
db.close()
def test_postgres_schema_is_current_idempotent_and_readers_are_read_only(tmp_path: Path):
db = Database(tmp_path / "schema-postgres")
db.initialize()
with db.transaction() as conn:
conn.executemany(
"INSERT INTO jobs(id,kind,payload_json,status,priority,run_after,created_at) VALUES(?,?,?,?,?,?,?)",
(
("legacy-audio-queued", "transcribe_audio", "{}", "queued", 30, "now", "now"),
("legacy-audio-complete", "transcribe_audio", "{}", "completed", 30, "now", "now"),
("legacy-index-queued", "index_video", "{}", "queued", 20, "now", "now"),
("legacy-index-complete", "index_video", "{}", "completed", 20, "now", "now"),
),
)
db.initialize()
assert db.setting("schema_version") == SCHEMA_VERSION
with db.read() as reader:
assert reader.execute("SELECT extversion FROM pg_extension WHERE extname='vector'").fetchone()[0]
upload_columns = {
row[0]
for row in reader.execute(
"SELECT column_name FROM information_schema.columns WHERE table_name='uploads'"
).fetchall()
}
assert {"title", "target_key", "webdav_path", "content_sha256_verified"} <= upload_columns
with db.read() as reader:
priorities = {
row["id"]: row["priority"]
for row in reader.execute(
"SELECT id,priority FROM jobs WHERE id IN ("
"'legacy-audio-queued','legacy-audio-complete','legacy-index-queued','legacy-index-complete')"
).fetchall()
}
assert reader.execute(
"SELECT 1 FROM pg_indexes WHERE indexname='idx_text_entries_raw_text_trgm'"
).fetchone()
assert priorities == {
"legacy-audio-queued": 20,
"legacy-audio-complete": 30,
"legacy-index-queued": 10,
"legacy-index-complete": 20,
}
with db.read() as reader:
reader_pid = reader.execute("SELECT pg_backend_pid()").fetchone()[0]
observer = db.connect()
try:
state = observer.execute(
"SELECT state FROM pg_stat_activity WHERE pid=?", (reader_pid,)
).fetchone()[0]
finally:
observer.rollback()
observer.close()
assert state == "idle"
with pytest.raises(Exception, match="read-only"):
with db.read() as reader:
reader.execute("INSERT INTO settings(key,value,updated_at) VALUES('forbidden','1','now')")
def test_schema_initialization_backfills_job_resources_with_legacy_postgres_json_syntax(tmp_path: Path):
db = Database(tmp_path / "job-resource-backfill-postgres")
db.initialize()
with db.transaction() as conn:
conn.executemany(
"INSERT INTO jobs(id,kind,payload_json,run_after,created_at) VALUES(?,?,?,?,?)",
(
("single-job", "index_video", '{"video_id":"video-one"}', "now", "now"),
(
"bulk-job",
"bulk_index",
'{"video_ids":["video-two","video-three","video-two"]}',
"now",
"now",
),
("unrelated-job", "scan_source", '{"source_id":"source-one"}', "now", "now"),
),
)
# Re-running initialize simulates an installation/upgrade which needs to
# populate the resource relation for jobs created before that table existed.
# The implementation deliberately uses jsonb casts/operators available in
# older supported PostgreSQL releases and must not require the PG16-only
# SQL/JSON ``IS JSON`` predicate.
db.initialize()
with db.read() as conn:
resources = {
(row["job_id"], row["resource_type"], row["resource_id"])
for row in conn.execute(
"SELECT job_id,resource_type,resource_id FROM job_resources ORDER BY job_id,resource_id"
).fetchall()
}
assert resources == {
("single-job", "video", "video-one"),
("bulk-job", "video", "video-two"),
("bulk-job", "video", "video-three"),
}
def test_gateway_media_tokens_follow_session_lifecycle(tmp_path: Path):
db = Database(tmp_path / "media-token-postgres")
db.initialize()
auth = AuthService(db, session_days=1)
auth.setup("gateway media token test password")
session_token, _, _ = auth.login_gateway("nas-admin", "Administrator")
media_token, expires = auth.create_gateway_media_token(session_token)
assert expires.isoformat()
assert auth.gateway_media_session(media_token, "nas-admin")["auth_source"] == "gateway"
assert auth.gateway_media_session(media_token, "another-admin") is None
auth.logout(session_token)
assert auth.gateway_media_session(media_token, "nas-admin") is None
with db.read() as conn:
assert conn.execute("SELECT count(*) FROM gateway_media_tokens").fetchone()[0] == 0
+408
View File
@@ -0,0 +1,408 @@
from __future__ import annotations
import json
import os
import re
import subprocess
import sys
import tomllib
from pathlib import Path
import pytest
from PIL import Image
PROJECT_ROOT = Path(__file__).resolve().parents[1]
FNOS_ROOT = PROJECT_ROOT / "fnos"
def _mock_app(tmp_path: Path) -> tuple[Path, Path, Path]:
app_root = tmp_path / "app"
runtime = app_root / "runtime"
(runtime / "wheels").mkdir(parents=True)
(runtime / "VERSION").write_text("0.4.3\n" + "a" * 64 + "\n")
(runtime / "runtime-core.txt").write_text("# mocked offline core\n")
(runtime / "imagefind-0.4.3-py3-none-any.whl").write_bytes(b"mock wheel")
(runtime / "wheels" / "dependency.whl").write_bytes(b"mock dependency")
python = app_root / "fake-python312"
python.write_text(
"""#!/bin/bash
set -eu
if [ "${1:-}" = "-c" ]; then
case "${2:-}" in
*imagefind.__version__*) printf '%s\n' '0.4.3'; exit 0 ;;
esac
if [ "$#" -gt 2 ]; then
exec "$MOCK_REAL_PYTHON" "$@"
fi
printf '%s\n' '3.12'
exit 0
fi
if [ "${1:-}" = "-m" ] && [ "${2:-}" = "venv" ]; then
mkdir -p "$3/bin"
cp "$0" "$3/bin/python"
chmod 0755 "$3/bin/python"
exit 0
fi
if [ "${1:-}" = "-m" ] && [ "${2:-}" = "pip" ]; then
exit 0
fi
if [ "${1:-}" = "-m" ] && [ "${2:-}" = "imagefind.main" ]; then
shift 2
fi
if [ "${1:-}" = "--version" ]; then
printf '%s\n' '0.4.3'
exit 0
fi
if [ "${1:-}" = "postgres-enroll" ]; then
rm -f "$TRIM_PKGVAR/postgres-enrollment-token.seed"
exit 0
fi
if [ "${1:-}" = "admin-password" ]; then
printf '%s' "$*" >"$MOCK_ARGS_FILE"
[ -z "${wizard_admin_password:-}" ]
[ -z "${wizard_admin_password_confirm:-}" ]
PASSWORD=''
IFS= read -r PASSWORD || true
printf '%s' "$PASSWORD" >"$MOCK_PASSWORD_FILE"
fi
""",
encoding="utf-8",
)
python.chmod(0o755)
return app_root, tmp_path / "mock-args", tmp_path / "mock-password"
def _environment(
tmp_path: Path,
app_root: Path,
args_file: Path,
password_file: Path,
) -> dict[str, str]:
return {
**os.environ,
"TRIM_APPDEST": str(app_root),
"TRIM_PKGVAR": str(tmp_path / "var"),
"TRIM_TEMP_LOGFILE": str(tmp_path / "fnos-error.log"),
"MOCK_ARGS_FILE": str(args_file),
"MOCK_PASSWORD_FILE": str(password_file),
"MOCK_REAL_PYTHON": sys.executable,
"IMAGEFIND_PYTHON_PATH": str(app_root / "fake-python312"),
}
def _run(script: str, env: dict[str, str]) -> subprocess.CompletedProcess[str]:
return subprocess.run(
["bash", str(FNOS_ROOT / "cmd" / script)],
env=env,
text=True,
capture_output=True,
check=False,
)
def test_fnos_wizard_contracts():
install = json.loads((FNOS_ROOT / "wizard" / "install").read_text())
upgrade = json.loads((FNOS_ROOT / "wizard" / "upgrade").read_text())
uninstall = json.loads((FNOS_ROOT / "wizard" / "uninstall").read_text())
config = json.loads((FNOS_ROOT / "wizard" / "config").read_text())
assert [item["type"] for item in install[0]["items"]] == ["tips", "password", "password", "password"]
assert install[0]["items"][1]["field"] == "wizard_postgres_enrollment_token"
assert [item["type"] for item in upgrade[0]["items"]] == ["tips"]
clear_data = uninstall[0]["items"][1]
assert clear_data["field"] == "wizard_clear_data"
# fnpack 1.2.3 requires switch defaults to be encoded as strings.
assert clear_data["initValue"] == "false"
assert [item["type"] for item in config[0]["items"]] == ["tips", "switch", "text"]
assert config[0]["items"][1]["initValue"] == "false"
desktop = json.loads((FNOS_ROOT / "app" / "ui" / "config").read_text())[".url"]["imagefind.Application"]
assert desktop["gatewayPrefix"] == "/app/imagefind"
assert desktop["gatewaySocket"] == "imagefind.sock"
assert desktop["url"] == "/app/imagefind/"
manifest = (FNOS_ROOT / "manifest").read_text()
assert "checkport=false" in manifest
assert "service_port=" not in manifest
manifest_version = re.search(r"^version\s*=\s*([^\s]+)$", manifest, re.MULTILINE).group(1)
project_version = tomllib.loads((PROJECT_ROOT / "pyproject.toml").read_text())["project"]["version"]
frontend_version = json.loads((PROJECT_ROOT / "frontend" / "package.json").read_text())["version"]
backend_version = re.search(
r'^__version__\s*=\s*"([^"]+)"$',
(PROJECT_ROOT / "backend" / "imagefind" / "__init__.py").read_text(),
re.MULTILINE,
).group(1)
assert {manifest_version, project_version, frontend_version, backend_version} == {"0.5.45"}
assert "install_dep_apps=python312,nxsir.postgresql" in manifest
assert os.access(FNOS_ROOT / "cmd" / "config_callback", os.X_OK)
assert os.access(FNOS_ROOT / "cmd" / "runtime_bootstrap", os.X_OK)
build_script = (PROJECT_ROOT / "scripts" / "build-fnos.sh").read_text()
assert "PyInstaller" not in build_script
assert "requirements/runtime-core.txt" in build_script
assert ".fnos-wheel-cache/python312" in build_script
assert "vendor/qdrant" not in build_script
assert "vendor/ffmpeg" not in build_script
assert f"imagefind-{manifest_version}-x86_64.fpk" in build_script
def test_generated_fnos_icons_fill_the_opaque_canvas(tmp_path: Path):
stage = tmp_path / "stage"
subprocess.run(
[sys.executable, str(PROJECT_ROOT / "scripts" / "make_icons.py"), str(stage)],
check=True,
)
expected = {
stage / "ICON.PNG": 128,
stage / "ICON_256.PNG": 256,
stage / "app" / "ui" / "images" / "icon_64.png": 64,
stage / "app" / "ui" / "images" / "icon_256.png": 256,
}
for path, size in expected.items():
with Image.open(path) as icon:
assert icon.mode == "RGBA"
assert icon.size == (size, size)
assert icon.getpixel((0, 0)) == (103, 145, 244, 255)
assert icon.getpixel((size - 1, size - 1)) == (103, 145, 244, 255)
assert icon.getpixel((size // 2, size // 2))[3] == 255
def test_install_initializes_password_without_logging_it(tmp_path: Path):
app_root, args_file, password_file = _mock_app(tmp_path)
env = _environment(tmp_path, app_root, args_file, password_file)
password = "administrator passphrase"
postgres_token = "test-postgresql-enrollment-token"
env.update(
wizard_postgres_enrollment_token=postgres_token,
wizard_admin_password=password,
wizard_admin_password_confirm=password,
)
result = _run("install_callback", env)
assert result.returncode == 0, result.stderr
assert args_file.read_text() == "admin-password --stdin --replace"
assert password_file.read_text() == password
assert (tmp_path / "var" / ".imagefind-data-root").read_text().strip() == "imagefind-data-v1"
assert json.loads((tmp_path / "var" / "access.json").read_text())["direct_access"] == {
"enabled": False,
"port": 8765,
}
assert password not in result.stdout
assert password not in result.stderr
assert password not in (tmp_path / "fnos-error.log").read_text()
assert postgres_token not in result.stdout
assert postgres_token not in result.stderr
assert postgres_token not in (tmp_path / "fnos-error.log").read_text()
assert not (tmp_path / "var" / "postgres-enrollment-token.seed").exists()
def test_install_rejects_mismatched_passwords_without_invoking_server(tmp_path: Path):
app_root, args_file, password_file = _mock_app(tmp_path)
env = _environment(tmp_path, app_root, args_file, password_file)
env.update(
wizard_postgres_enrollment_token="test-postgresql-enrollment-token",
wizard_admin_password="administrator passphrase",
wizard_admin_password_confirm="different administrator passphrase",
)
result = _run("install_callback", env)
assert result.returncode != 0
assert "不一致" in result.stderr
assert not args_file.exists()
assert not password_file.exists()
error_log = (tmp_path / "fnos-error.log").read_text()
assert "不一致" in error_log
assert env["wizard_admin_password"] not in error_log
assert env["wizard_admin_password_confirm"] not in error_log
def test_upgrade_is_idempotent_and_never_changes_password(tmp_path: Path):
app_root, args_file, password_file = _mock_app(tmp_path)
env = _environment(tmp_path, app_root, args_file, password_file)
data_root = tmp_path / "var"
(data_root / "data").mkdir(parents=True)
database = data_root / "data" / "imagefind.sqlite3"
database.write_text("existing database")
old_runtime = data_root / "runtime" / "ai-current"
old_runtime.mkdir(parents=True)
(old_runtime / "sentinel").write_text("preserve old runtime data")
(data_root / ".imagefind-data-root").write_text("imagefind-data-v1\n")
(data_root / "access.json").write_text('{"direct_access":{"enabled":false,"port":9876}}\n')
# Cached fnOS wizard values from an older package must not turn a file
# replacement into a failed update or reset the administrator password.
env.update(wizard_admin_password="placeholder", wizard_admin_password_confirm="different")
preserve = _run("upgrade_callback", env)
again = _run("upgrade_callback", env)
assert preserve.returncode == 0, preserve.stderr
assert again.returncode == 0, again.stderr
assert not args_file.exists()
assert not password_file.exists()
assert database.read_text() == "existing database"
assert (old_runtime / "sentinel").read_text() == "preserve old runtime data"
assert json.loads((data_root / "access.json").read_text())["direct_access"] == {
"enabled": False,
"port": 9876,
}
assert "升级数据检查完成" in (tmp_path / "fnos-error.log").read_text()
def test_upgrade_init_does_not_migrate_or_modify_legacy_sqlite_data(tmp_path: Path):
app_root = tmp_path / "app"
(app_root / "bin").mkdir(parents=True)
qdrant = app_root / "bin" / "qdrant"
qdrant.write_bytes(b"legacy-qdrant-binary")
qdrant.chmod(0o755)
(app_root / "qdrant.yaml").write_text("storage:\n storage_path: ./storage\n")
data_root = tmp_path / "var"
database = data_root / "data" / "imagefind.sqlite3"
database.parent.mkdir(parents=True)
database.write_bytes(b"legacy sqlite data remains outside the PostgreSQL release")
env = {
**os.environ,
"TRIM_APPDEST": str(app_root),
"TRIM_PKGVAR": str(data_root),
"TRIM_TEMP_LOGFILE": str(tmp_path / "fnos-upgrade.log"),
"IMAGEFIND_PYTHON_PATH": sys.executable,
}
result = _run("upgrade_init", env)
assert result.returncode == 0, result.stderr
assert database.read_bytes() == b"legacy sqlite data remains outside the PostgreSQL release"
assert qdrant.read_bytes() == b"legacy-qdrant-binary"
assert not (data_root / "migration").exists()
assert "不执行旧数据迁移" in (tmp_path / "fnos-upgrade.log").read_text()
@pytest.mark.parametrize(
("current_kind", "old_name"),
[
("symlink", "core-0.3.14"),
("symlink", "core-0.3.16"),
("dangling", "core-missing"),
("directory", "current"),
],
)
def test_runtime_bootstrap_atomically_replaces_existing_current(
tmp_path: Path,
current_kind: str,
old_name: str,
):
app_root, args_file, password_file = _mock_app(tmp_path)
env = _environment(tmp_path, app_root, args_file, password_file)
runtime_root = tmp_path / "var" / "runtime"
runtime_root.mkdir(parents=True)
current = runtime_root / "current"
if current_kind == "symlink":
old_runtime = runtime_root / old_name
old_runtime.mkdir()
(old_runtime / "sentinel").write_text("preserve old runtime")
current.symlink_to(old_name)
elif current_kind == "dangling":
old_runtime = runtime_root / old_name
current.symlink_to(old_name)
else:
old_runtime = current
old_runtime.mkdir()
(old_runtime / "sentinel").write_text("preserve old runtime")
result = _run("runtime_bootstrap", env)
assert result.returncode == 0, result.stderr
assert current.is_symlink()
assert current.resolve() == runtime_root / "core-0.4.3"
version = subprocess.run(
[str(current / "bin" / "python"), "-m", "imagefind.main", "--version"],
env=env,
text=True,
capture_output=True,
check=True,
)
assert version.stdout.strip() == "0.4.3"
if current_kind == "symlink":
assert (old_runtime / "sentinel").read_text() == "preserve old runtime"
assert not list(old_runtime.glob(".current.*"))
elif current_kind == "dangling":
assert not old_runtime.exists()
else:
previous = list(runtime_root.glob("current.previous.*"))
assert len(previous) == 1
assert (previous[0] / "sentinel").read_text() == "preserve old runtime"
def test_config_callback_writes_validated_direct_access_atomically(tmp_path: Path):
data_root = tmp_path / "var"
env = {
**os.environ,
"TRIM_PKGVAR": str(data_root),
"TRIM_TEMP_LOGFILE": str(tmp_path / "fnos-error.log"),
"wizard_direct_access": "true",
"wizard_direct_port": "9876",
}
configured = _run("config_callback", env)
assert configured.returncode == 0, configured.stderr
assert json.loads((data_root / "access.json").read_text()) == {
"direct_access": {"enabled": True, "port": 9876}
}
assert (data_root / "access.json").stat().st_mode & 0o777 == 0o600
assert not list(data_root.glob("access.json.tmp.*"))
env["wizard_direct_port"] = "70000"
invalid = _run("config_callback", env)
assert invalid.returncode != 0
assert json.loads((data_root / "access.json").read_text())["direct_access"]["port"] == 9876
def test_uninstall_preserves_by_default_and_only_clears_internal_data(tmp_path: Path):
data_root = tmp_path / "var"
data_root.mkdir()
(data_root / ".imagefind-data-root").write_text("imagefind-data-v1\n")
(data_root / "data").mkdir()
(data_root / "data" / "imagefind.sqlite3").write_text("database")
(data_root / "models").mkdir()
(data_root / "models" / "model.bin").write_text("model")
shared_root = tmp_path / "shared" / "imagefind" / "videos"
shared_root.mkdir(parents=True)
source_video = shared_root / "source.mp4"
source_video.write_bytes(b"source-video")
app_root = tmp_path / "app"
env = {
**os.environ,
"TRIM_APPDEST": str(app_root),
"TRIM_PKGVAR": str(data_root),
"TRIM_TEMP_LOGFILE": str(tmp_path / "fnos-error.log"),
"TRIM_DATA_SHARE_PATHS": str(shared_root),
"TRIM_DATA_ACCESSIBLE_PATHS": str(tmp_path / "authorized-videos"),
}
preserve = _run("uninstall_init", env)
assert preserve.returncode == 0, preserve.stderr
assert (data_root / "data" / "imagefind.sqlite3").exists()
assert source_video.read_bytes() == b"source-video"
env["wizard_clear_data"] = "true"
clear = _run("uninstall_init", env)
assert clear.returncode == 0, clear.stderr
assert list(data_root.iterdir()) == []
assert source_video.read_bytes() == b"source-video"
def test_uninstall_refuses_data_directory_without_safety_marker(tmp_path: Path):
data_root = tmp_path / "not-imagefind-data"
data_root.mkdir()
protected = data_root / "keep.txt"
protected.write_text("keep")
env = {
**os.environ,
"TRIM_PKGVAR": str(data_root),
"TRIM_TEMP_LOGFILE": str(tmp_path / "fnos-error.log"),
"wizard_clear_data": "true",
}
result = _run("uninstall_init", env)
assert result.returncode != 0
assert "安全标记" in result.stderr
assert protected.read_text() == "keep"
+6
View File
@@ -0,0 +1,6 @@
from imagefind.indexer import Indexer
def test_text_replacement_uses_portable_fixed_predicates():
assert Indexer._replaceable_text_condition(audio_stale=True) == ""
assert Indexer._replaceable_text_condition(audio_stale=False) == " AND kind<>'audio'"
+41
View File
@@ -0,0 +1,41 @@
from __future__ import annotations
import time
from pathlib import Path
from imagefind.accelerator import AcceleratorService
from imagefind.config import Settings
from imagefind.inference import InferenceSupervisor, IsolatedEmbeddingService
def test_hash_embeddings_remain_in_process_without_starting_worker(tmp_path: Path):
settings = Settings(data_dir=tmp_path, embedding_backend="hash")
settings.prepare()
accelerator = AcceleratorService(settings)
supervisor = InferenceSupervisor(settings, accelerator, idle_seconds=5)
embeddings = IsolatedEmbeddingService(settings, accelerator, supervisor)
vector = embeddings.encode_text("本地测试")
assert len(vector) == 512
assert supervisor.status()["running"] is False
supervisor.close()
def test_inference_worker_is_lazy_and_exits_after_idle_window(tmp_path: Path):
settings = Settings(data_dir=tmp_path, embedding_backend="auto")
settings.prepare()
accelerator = AcceleratorService(settings)
supervisor = InferenceSupervisor(settings, accelerator, idle_seconds=5)
try:
assert supervisor.status()["running"] is False
assert supervisor.call("reset", component="visual") is None
started = supervisor.status()
assert started["running"] is True
assert started["pid"]
deadline = time.monotonic() + 8
while supervisor.status()["running"] and time.monotonic() < deadline:
time.sleep(0.1)
assert supervisor.status()["running"] is False
finally:
supervisor.close()
+483
View File
@@ -0,0 +1,483 @@
from __future__ import annotations
import os
from pathlib import Path
import httpx
from imagefind.config import Settings
from imagefind.remote import RcloneManager
from imagefind.remote_cache import RemoteMediaCache
from imagefind.resources import ResourceGovernor
from imagefind.usage import StorageUsageService
class _SettingsDb:
def __init__(self):
self.values = {}
def setting(self, key, default=None):
return self.values.get(key, default)
def set_setting(self, key, value):
self.values[key] = value
class _RemoteSource:
def remote_access(self, _source_id: str, _key: str):
return "https://media.test/video.mp4", "user", "password", True
def _sample(some: float, full: float) -> dict:
return {
"cpu_percent": 0.0,
"memory_total_bytes": 8 * 1024**3,
"memory_available_bytes": 6 * 1024**3,
"disk_available_bytes": 100 * 1024**3,
"io": {
"supported": True,
"some_avg10": some,
"some_avg60": some,
"some_avg300": some,
"full_avg10": full,
"full_avg60": full,
"full_avg300": full,
},
}
def test_io_pressure_profiles_use_hysteresis(tmp_path: Path):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
db = _SettingsDb()
governor = ResourceGovernor(db, settings)
governor.sample = lambda: _sample(31, 11) # type: ignore[method-assign]
for generation in range(1, 4):
governor._sample_generation = generation
state = governor._update_io_state(governor.sample())
assert state == "paused"
assert "磁盘 I/O" in governor.pressure_reason(lane="ai", running=True)
governor.sample = lambda: _sample(2, 0) # type: ignore[method-assign]
for generation in range(4, 9):
governor._sample_generation = generation
state = governor._update_io_state(governor.sample())
assert state == "normal"
db.set_setting("resource_profile", "quiet")
assert governor.profile() == "quiet"
def test_io_pressure_uses_data_volume_instead_of_system_psi(tmp_path: Path):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
governor = ResourceGovernor(_SettingsDb(), settings)
sample = _sample(95, 80)
sample["io"].update(
{
"scope": "data_volume",
"volume_supported": True,
"volume_sampled": True,
"volume_utilization_percent": 4.0,
"volume_in_flight": 0,
}
)
for generation in range(1, 5):
governor._sample_generation = generation
state = governor._update_io_state(sample)
assert state == "normal"
sample["io"]["volume_utilization_percent"] = 92.0
for generation in range(5, 8):
governor._sample_generation = generation
state = governor._update_io_state(sample)
assert state == "paused"
def test_data_volume_block_device_is_resolved_from_mountinfo(tmp_path: Path):
data = tmp_path / "volume" / "app" / "data"
data.mkdir(parents=True)
mountinfo = tmp_path / "mountinfo"
mountinfo.write_text(
f"10 1 8:1 / {tmp_path / 'volume'} rw - ext4 /dev/test rw\n",
encoding="utf-8",
)
sysfs = tmp_path / "sys" / "dev" / "block"
stat = sysfs / "8:1" / "stat"
stat.parent.mkdir(parents=True)
stat.write_text("1 0 0 0 2 0 0 0 3 456 0", encoding="utf-8")
result = ResourceGovernor._data_volume_io_values(
data,
mountinfo_path=mountinfo,
sysfs_root=sysfs,
)
assert result == {
"device": "8:1",
"mount_point": str(tmp_path / "volume"),
"in_flight": 3,
"io_ticks_ms": 456,
}
def test_remote_media_is_downloaded_once_and_reused(tmp_path: Path, monkeypatch):
payload = b"0123456789" * 1024
requests = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
range_header = request.headers.get("range")
if range_header:
offset = int(range_header.removeprefix("bytes=").removesuffix("-"))
return httpx.Response(
206,
headers={"Content-Range": f"bytes {offset}-{len(payload) - 1}/{len(payload)}"},
content=payload[offset:],
)
return httpx.Response(200, content=payload)
transport = httpx.MockTransport(handler)
client_type = httpx.Client
monkeypatch.setattr(
"imagefind.remote_cache.httpx.Client",
lambda **kwargs: client_type(transport=transport, follow_redirects=True),
)
settings = Settings(
data_dir=tmp_path / "data",
remote_cache_gb=0.01,
resource_disk_reserve_gb=0.5,
)
settings.prepare()
cache = RemoteMediaCache(settings, _RemoteSource())
video = {
"id": "video-1",
"source_id": "source-1",
"source_key": "folder/video.mp4",
"fingerprint": "etag-1",
"size_bytes": len(payload),
}
first = cache.acquire(video)
assert first.path.read_bytes() == payload
first.close()
second = cache.acquire(video)
assert second.path == first.path
second.close()
assert len(requests) == 1
assert cache.status()["used_bytes"] == len(payload)
def test_remote_media_partial_download_resumes(tmp_path: Path, monkeypatch):
payload = b"abcdefghij"
observed_range = []
def handler(request: httpx.Request) -> httpx.Response:
range_header = request.headers.get("range")
observed_range.append(range_header)
offset = int((range_header or "bytes=0-").removeprefix("bytes=").removesuffix("-"))
return httpx.Response(
206 if offset else 200,
headers={"Content-Range": f"bytes {offset}-{len(payload) - 1}/{len(payload)}"},
content=payload[offset:],
)
transport = httpx.MockTransport(handler)
client_type = httpx.Client
monkeypatch.setattr(
"imagefind.remote_cache.httpx.Client",
lambda **kwargs: client_type(transport=transport, follow_redirects=True),
)
settings = Settings(data_dir=tmp_path / "data", resource_disk_reserve_gb=0.5)
settings.prepare()
cache = RemoteMediaCache(settings, _RemoteSource())
video = {
"id": "video-2",
"source_id": "source-1",
"source_key": "video.mp4",
"fingerprint": "etag-2",
"size_bytes": len(payload),
}
_key, _data, part, _metadata = cache._paths(video)
part.write_bytes(payload[:4])
lease = cache.acquire(video)
assert lease.path.read_bytes() == payload
lease.close()
assert observed_range == ["bytes=4-"]
def test_active_partial_is_pinned_against_manual_cleanup(tmp_path: Path, monkeypatch):
payload = b"active-download"
settings = Settings(data_dir=tmp_path / "data", resource_disk_reserve_gb=0.5)
settings.prepare()
cache = RemoteMediaCache(settings, _RemoteSource())
video = {
"id": "video-active",
"source_id": "source-1",
"source_key": "active.mp4",
"fingerprint": "etag-active",
"size_bytes": len(payload),
}
def download(_video, part, **_kwargs):
part.write_bytes(payload)
result = cache.cleanup(force=True)
assert result["skipped_active"] == 1
assert part.read_bytes() == payload
monkeypatch.setattr(cache, "_download", download)
lease = cache.acquire(video)
assert lease.path.read_bytes() == payload
lease.close()
def test_rclone_instances_are_isolated_and_removed_on_stop(tmp_path: Path, monkeypatch):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
manager = RcloneManager(settings)
commands = []
ports = iter((41001, 41002))
class Process:
next_pid = 9100
def __init__(self):
self.pid = Process.next_pid
Process.next_pid += 1
self.returncode = None
def poll(self):
return self.returncode
def terminate(self):
self.returncode = 0
def kill(self):
self.returncode = -9
def wait(self, timeout=None):
return self.returncode
class Connection:
def __enter__(self):
return self
def __exit__(self, *_args):
return None
def popen(command, **_kwargs):
commands.append(command)
return Process()
monkeypatch.setattr(manager, "_binary", lambda: "/usr/bin/rclone")
monkeypatch.setattr(manager, "_obscure", lambda value: f"obscured-{value}")
monkeypatch.setattr(manager, "_free_port", lambda: next(ports))
monkeypatch.setattr("imagefind.remote.subprocess.Popen", popen)
monkeypatch.setattr("imagefind.remote.socket.create_connection", lambda *_args, **_kwargs: Connection())
def source(source_id):
return {
"id": source_id,
"config": {
"driver": "alist",
"mode": "encrypted",
"base_url": "https://alist.test",
"root_path": "videos",
"username": "admin",
},
"secrets": {"password": "dav", "crypt_password": "crypt", "crypt_salt": "salt"},
}
first = manager.ensure(source("source-a"))
second = manager.ensure(source("source-b"))
first_cache = Path(commands[0][commands[0].index("--cache-dir") + 1])
second_cache = Path(commands[1][commands[1].index("--cache-dir") + 1])
assert first_cache != second_cache
assert first.instance_dir in first_cache.parents
assert second.instance_dir in second_cache.parents
assert "--read-only" in commands[0]
assert commands[0][commands[0].index("--vfs-cache-mode") + 1] == "off"
assert "--vfs-cache-max-size" not in commands[0]
assert commands[0][commands[0].index("--buffer-size") + 1] == "4M"
manager.stop_all()
assert not first.instance_dir.exists()
assert not second.instance_dir.exists()
def test_rclone_start_failure_does_not_leave_instance_files(tmp_path: Path, monkeypatch):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
manager = RcloneManager(settings)
monkeypatch.setattr(manager, "_binary", lambda: "/usr/bin/rclone")
monkeypatch.setattr(manager, "_obscure", lambda value: value)
monkeypatch.setattr(manager, "_free_port", lambda: 41003)
def fail_start(*_args, **_kwargs):
raise OSError("cannot start")
monkeypatch.setattr("imagefind.remote.subprocess.Popen", fail_start)
source = {
"id": "source-failed",
"config": {
"driver": "alist",
"mode": "encrypted",
"base_url": "https://alist.test",
"root_path": "",
"username": "admin",
},
"secrets": {"password": "dav", "crypt_password": "crypt", "crypt_salt": "salt"},
}
try:
manager.ensure(source)
except OSError as exc:
assert "cannot start" in str(exc)
else:
raise AssertionError("rclone startup failure must be surfaced")
assert not list((settings.rclone_dir / "instances" / "source-failed").glob("*"))
def test_rclone_orphan_is_reclaimed_after_three_idle_samples(tmp_path: Path, monkeypatch):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
reconcile = RcloneManager.reconcile_orphans
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
manager = RcloneManager(settings)
monkeypatch.setattr(RcloneManager, "reconcile_orphans", reconcile)
target_pid = os.getpid()
times = iter((1000.0, 1060.0, 1121.0, 1122.0))
reaped: list[int] = []
monkeypatch.setattr("imagefind.remote.time.monotonic", lambda: next(times))
monkeypatch.setattr(manager, "_process_identity", lambda pid: (1, 77) if pid == target_pid else None)
monkeypatch.setattr(manager, "_managed_command", lambda pid: pid == target_pid)
monkeypatch.setattr(manager, "_registry_matches", lambda _pid, _ticks: None)
monkeypatch.setattr(manager, "_io_bytes", lambda _pid: (10, 20))
monkeypatch.setattr(manager, "_has_established_connection", lambda _pid: False)
monkeypatch.setattr(manager, "_has_open_cache_file", lambda _pid: True)
monkeypatch.setattr(manager, "_terminate_pid", lambda pid: reaped.append(pid))
monkeypatch.setattr(manager, "_cleanup_registered_instance", lambda *_args: None)
assert manager.reconcile_orphans()["verifying"] == 1
assert manager.reconcile_orphans()["verifying"] == 1
result = manager.reconcile_orphans()
assert result == {"matched": 1, "reaped": 1, "verifying": 0, "unresolved": 0}
assert reaped == [target_pid]
assert manager.status()["reclaimed"] == 1
def test_rclone_cleanup_removes_only_stale_vfs_cache(tmp_path: Path, monkeypatch):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
manager = RcloneManager(settings)
stale = settings.rclone_dir / "instances" / "source-stale" / "instance-stale"
cache = stale / "cache" / "vfs" / "encrypted"
cache.mkdir(parents=True)
(cache / "movie.bin").write_bytes(b"cached" * 1024)
(stale / "instance.json").write_text('{"pid":999999,"start_ticks":1}', encoding="utf-8")
before = manager.status()
assert before["vfs_cache_bytes"] == 6 * 1024
assert before["reclaimable_cache_bytes"] == 6 * 1024
result = manager.cleanup_vfs_cache()
assert result == {"freed_bytes": 6 * 1024, "removed": 1, "skipped_active": 0}
assert not (stale / "cache").exists()
def test_rclone_accounts_for_and_cleans_legacy_shared_vfs_cache(tmp_path: Path, monkeypatch):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
manager = RcloneManager(settings)
legacy = settings.rclone_dir / "cache" / "vfs" / "crypt_source"
legacy.mkdir(parents=True)
(legacy / "old-video.bin").write_bytes(b"legacy" * 1024)
before = manager.status()
assert before["vfs_cache_bytes"] == 6 * 1024
assert before["reclaimable_cache_bytes"] == 6 * 1024
assert before["runtime_bytes"] == 0
assert before["instances"][0]["instance_id"] == "legacy-cache"
result = manager.cleanup_vfs_cache()
assert result == {"freed_bytes": 6 * 1024, "removed": 1, "skipped_active": 0}
assert not (settings.rclone_dir / "cache").exists()
def test_storage_usage_exposes_nested_cache_and_staging_details(tmp_path: Path):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
(settings.models_dir / "model.bin").write_bytes(b"m" * 10)
(settings.remote_media_cache_dir / "video.media").write_bytes(b"r" * 20)
external = tmp_path / "openlist-stage"
external.mkdir()
(external / "encrypted.bin").write_bytes(b"e" * 30)
outside = tmp_path / "outside.bin"
outside.write_bytes(b"x" * 100)
(settings.remote_media_cache_dir / "ignored-link").symlink_to(outside)
class Connection:
def execute(self, sql, _params=()):
if "pg_database_size" in sql:
return type("Result", (), {"fetchone": lambda self: (1000,)})()
if "pg_total_relation_size" in sql:
return type("Result", (), {"fetchone": lambda self: (200,)})()
if "FROM sources" in sql:
rows = [
{
"id": "source-1",
"name": "OpenList",
"config_json": f'{{"openlist_local_staging_path":"{external}"}}',
}
]
return type("Result", (), {"fetchall": lambda self: rows})()
raise AssertionError(sql)
class Read:
def __enter__(self):
return Connection()
def __exit__(self, *_args):
return None
class Db:
def read(self):
return Read()
class Rclone:
def status(self):
return {
"vfs_cache_bytes": 40,
"active_cache_bytes": 5,
"reclaimable_cache_bytes": 35,
"runtime_bytes": 3,
"instances": [],
}
class Cache:
def status(self):
return {"active_bytes": 4, "evictable_bytes": 16, "entries": 1}
usage = StorageUsageService(Db(), settings, Rclone(), Cache()).usage(refresh=True)
assert usage["schema_version"] == 2
assert usage["categories"]["remote_cache"]["children"]["remote_media_cache"]["bytes"] == 20
assert usage["categories"]["remote_cache"]["children"]["rclone_vfs_cache"]["bytes"] == 40
openlist = usage["categories"]["upload_staging"]["children"]["openlist_local_staging"]
assert openlist["bytes"] == 30
assert openlist["available"] is True
assert usage["app_bytes"] == _path_bytes(settings.data_dir) + 1000 + 30
def _path_bytes(root: Path) -> int:
return sum(path.stat().st_size for path in root.rglob("*") if path.is_file() and not path.is_symlink())
+35
View File
@@ -0,0 +1,35 @@
import hashlib
import json
import runpy
import sys
import tarfile
from pathlib import Path
def test_model_bundle_embeds_manifest_without_mutating_source(tmp_path: Path, monkeypatch):
source = tmp_path / "source"
image_dir = source / "visual" / "image"
text_dir = source / "visual" / "text"
image_dir.mkdir(parents=True)
text_dir.mkdir(parents=True)
model = image_dir / "model.bin"
model.write_bytes(b"image model")
(text_dir / "config.json").write_text("{}", encoding="utf-8")
source_manifest = source / "manifest.json"
source_manifest.write_text("source sentinel", encoding="utf-8")
output = tmp_path / "models.tar.gz"
script = Path(__file__).parents[1] / "scripts" / "build-model-bundle.py"
namespace = runpy.run_path(str(script))
monkeypatch.setattr(sys, "argv", [str(script), str(source), str(output), "--version", "test-v1"])
namespace["main"]()
assert source_manifest.read_text(encoding="utf-8") == "source sentinel"
with tarfile.open(output, "r:gz") as archive:
assert archive.getnames().count("manifest.json") == 1
manifest_file = archive.extractfile("manifest.json")
assert manifest_file is not None
manifest = json.load(manifest_file)
assert manifest["version"] == "test-v1"
assert "manifest.json" not in manifest["files"]
assert manifest["files"]["visual/image/model.bin"] == hashlib.sha256(model.read_bytes()).hexdigest()
+535
View File
@@ -0,0 +1,535 @@
from __future__ import annotations
import asyncio
import hashlib
import io
import json
import sys
import tarfile
import types
from pathlib import Path
import httpx
import pytest
from imagefind.config import Settings
from imagefind.main import create_app
from imagefind.models import ModelManager
from imagefind.speech import SpeechService
def _bundle(*, unsafe: bool = False) -> bytes:
output = io.BytesIO()
with tarfile.open(fileobj=output, mode="w:gz") as archive:
files = {
"manifest.json": json.dumps({"version": "manual-test"}).encode(),
"visual/image/modules.json": json.dumps([{"idx": 0, "path": "", "type": "Transformer"}]).encode(),
"visual/text/modules.json": json.dumps([{"idx": 0, "path": "", "type": "Transformer"}]).encode(),
}
if unsafe:
files["../outside"] = b"unsafe"
for name, content in files.items():
member = tarfile.TarInfo(name)
member.size = len(content)
archive.addfile(member, io.BytesIO(content))
return output.getvalue()
def _app(tmp_path: Path):
settings = Settings(
data_dir=tmp_path / "data",
embedding_backend="hash",
upload_reserve_gb=0,
model_upload_gb=1,
)
settings.prepare()
app = create_app(settings)
_, token = app.state.services.auth.create_api_token("models")
return app, {"Authorization": f"Bearer {token}"}
def _write_audio_variant(root: Path) -> None:
root.mkdir(parents=True, exist_ok=True)
for name in (
"config.json",
"preprocessor_config.json",
"tokenizer_config.json",
"tokenizer.json",
"openvino_encoder_model.xml",
"openvino_encoder_model.bin",
"openvino_decoder_model.xml",
"openvino_decoder_model.bin",
):
(root / name).write_text("{}", encoding="utf-8")
def test_manual_model_bundle_upload_and_atomic_install(tmp_path: Path):
app, headers = _app(tmp_path)
bundle = _bundle()
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/upload",
headers={
**headers,
"Content-Type": "application/octet-stream",
"X-Model-Filename": "imagefind-models.tar.gz",
},
content=bundle,
)
assert response.status_code == 202
payload = response.json()
assert payload["size_bytes"] == len(bundle)
assert payload["sha256"] == hashlib.sha256(bundle).hexdigest()
with app.state.services.db.read() as conn:
row = conn.execute("SELECT payload_json FROM jobs WHERE id=?", (payload["job_id"],)).fetchone()
job_payload = json.loads(row["payload_json"])
archive_path = Path(job_payload["archive_path"])
result = app.state.services.models.install(
archive_path=archive_path,
expected_sha256=job_payload["sha256"],
)
assert result["manifest"]["version"] == "manual-test"
assert (app.state.services.settings.models_dir / "visual" / "image" / "modules.json").is_file()
assert (app.state.services.settings.models_dir / "visual" / "text" / "modules.json").is_file()
asyncio.run(scenario())
def test_medium_audio_variant_uses_isolated_verification(tmp_path: Path, monkeypatch):
app, _headers = _app(tmp_path)
service = app.state.services
_write_audio_variant(service.settings.models_dir / "audio" / "small")
_write_audio_variant(service.settings.models_dir / "audio" / "medium")
service.settings.audio_model_variant = "small"
service.speech.accelerator.gpu_device = "GPU.0"
called: list[str] = []
def forbidden_main_process_verify():
raise AssertionError("Medium verification must not run in the API process")
def isolated_verify(variant):
called.append(f"{variant}:{service.settings.audio_model_variant}")
service.speech.accelerator.mark_ready("audio", "GPU.0", ["GPU.0"])
return service.speech.accelerator.status()["components"]["audio"]
monkeypatch.setattr(service.speech, "verify_acceleration", forbidden_main_process_verify)
monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)
selected = service.models.set_audio_variant("medium")
assert selected["active"] is True
assert called == ["medium:small"]
assert service.settings.audio_model_variant == "medium"
def test_medium_audio_variant_verification_failure_rolls_back(tmp_path: Path, monkeypatch):
app, _headers = _app(tmp_path)
service = app.state.services
_write_audio_variant(service.settings.models_dir / "audio" / "small")
_write_audio_variant(service.settings.models_dir / "audio" / "medium")
service.settings.audio_model_variant = "small"
service.speech.accelerator.gpu_device = "GPU.0"
def isolated_verify(_variant):
raise RuntimeError("worker crashed")
monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)
with pytest.raises(RuntimeError, match="worker crashed"):
service.models.set_audio_variant("medium")
assert service.settings.audio_model_variant == "small"
def test_speech_config_medium_failure_returns_conflict_and_keeps_small(tmp_path: Path, monkeypatch):
app, headers = _app(tmp_path)
service = app.state.services
_write_audio_variant(service.settings.models_dir / "audio" / "small")
_write_audio_variant(service.settings.models_dir / "audio" / "medium")
service.settings.audio_model_variant = "small"
service.speech.accelerator.gpu_device = "GPU.0"
def isolated_verify(_variant):
raise RuntimeError("worker crashed")
monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.patch(
"/api/v1/speech/config",
headers=headers,
json={"model_variant": "medium"},
)
assert response.status_code == 409
config = await client.get("/api/v1/speech/config", headers=headers)
assert config.json()["model_variant"] == "small"
asyncio.run(scenario())
def test_old_medium_fp16_variant_requires_int8_update(tmp_path: Path):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
medium = settings.models_dir / "audio" / "medium"
_write_audio_variant(medium)
manifest = {
"components": {
"audio": {
"variants": {
"medium": {
"version": "OpenVINO/whisper-medium-fp16-ov",
"source": "https://hf-mirror.com/OpenVINO/whisper-medium-fp16-ov",
}
},
"active_variant": "small",
}
}
}
(settings.models_dir / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
manager = ModelManager.__new__(ModelManager)
manager.settings = settings
manager.speech = SpeechService(settings)
manager._manifest_cached_mtime_ns = -1
manager._manifest_cached = {}
variant = manager.audio_variants()["medium"]
assert variant["installed"] is False
assert variant["needs_update"] is True
assert variant["expected_version"] == "OpenVINO/whisper-medium-int8-ov"
assert variant["size_bytes"] > 0
def test_model_status_defaults_to_light_manifest_and_allows_full_manifest(tmp_path: Path):
app, headers = _app(tmp_path)
app.state.services.settings.models_dir.mkdir(parents=True, exist_ok=True)
(app.state.services.settings.models_dir / "manifest.json").write_text(
json.dumps({"version": "test", "components": {}, "files": {"audio/model.bin": "hash"}}),
encoding="utf-8",
)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
light = await client.get("/api/v1/models", headers=headers)
assert light.status_code == 200
assert "files" not in light.json()["manifest"]
full = await client.get("/api/v1/models?summary=false", headers=headers)
assert full.status_code == 200
assert full.json()["manifest"]["files"] == {"audio/model.bin": "hash"}
asyncio.run(scenario())
def test_manual_model_bundle_rejects_unsafe_archive(tmp_path: Path):
app, _ = _app(tmp_path)
archive_path = tmp_path / "unsafe.tar.gz"
archive_path.write_bytes(_bundle(unsafe=True))
try:
app.state.services.models.install(archive_path=archive_path)
except ValueError as exc:
assert "不安全路径" in str(exc)
else:
raise AssertionError("unsafe archive must be rejected")
assert not (tmp_path / "outside").exists()
def test_default_model_install_uses_official_repositories(tmp_path: Path, monkeypatch):
app, _ = _app(tmp_path)
manager = app.state.services.models
rapid_root = tmp_path / "rapidocr"
rapid_models = rapid_root / "models"
rapid_models.mkdir(parents=True)
for name in (
"ch_PP-OCRv4_det_infer.onnx",
"ch_PP-OCRv4_rec_infer.onnx",
"ch_ppocr_mobile_v2.0_cls_infer.onnx",
):
(rapid_models / name).write_bytes(name.encode())
rapid_module = types.ModuleType("rapidocr_onnxruntime")
rapid_module.__file__ = str(rapid_root / "__init__.py")
hub_module = types.ModuleType("huggingface_hub")
requests_module = types.ModuleType("requests")
class Session:
def __init__(self):
self.trust_env = True
self.proxies = {}
def snapshot_download(repo_id, *, local_dir, **_):
target = Path(local_dir)
target.mkdir(parents=True, exist_ok=True)
(target / "modules.json").write_text(
json.dumps([{"idx": 0, "path": "", "type": "Transformer", "repo": repo_id}]),
encoding="utf-8",
)
return str(target)
hub_module.snapshot_download = snapshot_download
hub_module.configure_http_backend = lambda **_: None
requests_module.Session = Session
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", rapid_module)
monkeypatch.setitem(sys.modules, "huggingface_hub", hub_module)
monkeypatch.setitem(sys.modules, "requests", requests_module)
def download_file(url, destination):
destination.parent.mkdir(parents=True, exist_ok=True)
if destination.suffix == ".xml":
destination.write_text("<net><layers/></net>", encoding="utf-8")
else:
destination.write_bytes(b"\0" * (64 * 1024))
monkeypatch.setattr(manager, "_download_file", download_file)
def export_audio_model(_source, destination):
destination.mkdir(parents=True, exist_ok=True)
for name in (
"config.json",
"preprocessor_config.json",
"tokenizer_config.json",
"vocab.json",
"openvino_encoder_model.xml",
"openvino_encoder_model.bin",
"openvino_decoder_model.xml",
"openvino_decoder_model.bin",
"openvino_decoder_with_past_model.xml",
"openvino_decoder_with_past_model.bin",
):
(destination / name).write_text("{}" if name.endswith(".json") else "model")
monkeypatch.setattr(manager, "_export_audio_model", export_audio_model)
monkeypatch.setattr(manager, "repair_component", lambda *_args, **_kwargs: {})
progress = []
result = manager.install(progress=lambda value, message: progress.append((value, message)))
assert result["manifest"]["source"] == "official-repositories"
assert (manager.settings.models_dir / "ocr" / "det.onnx").is_file()
assert (manager.settings.models_dir / "faces" / "detector.xml").is_file()
assert (manager.settings.models_dir / "audio" / "openvino_encoder_model.xml").is_file()
assert any("多语言文本模型" in message for _, message in progress)
def test_component_install_preserves_other_components_and_manifest(tmp_path: Path, monkeypatch):
app, _ = _app(tmp_path)
manager = app.state.services.models
root = manager.settings.models_dir
(root / "faces").mkdir(parents=True)
(root / "faces" / "sentinel.bin").write_bytes(b"existing-face-index-runtime")
(root / "manifest.json").write_text(
json.dumps(
{
"format_version": 2,
"version": "existing",
"components": {"faces": {"version": "faces-existing"}},
"files": {"faces/sentinel.bin": hashlib.sha256(b"existing-face-index-runtime").hexdigest()},
}
),
encoding="utf-8",
)
rapid_root = tmp_path / "component-rapidocr"
rapid_models = rapid_root / "models"
rapid_models.mkdir(parents=True)
for name in (
"ch_PP-OCRv4_det_infer.onnx",
"ch_PP-OCRv4_rec_infer.onnx",
"ch_ppocr_mobile_v2.0_cls_infer.onnx",
):
(rapid_models / name).write_bytes(name.encode())
rapid_module = types.ModuleType("rapidocr_onnxruntime")
rapid_module.__file__ = str(rapid_root / "__init__.py")
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", rapid_module)
progress: list[tuple[float, str]] = []
result = manager.install(component="ocr", progress=lambda value, message: progress.append((value, message)))
assert result["component"] == "ocr"
assert (root / "faces" / "sentinel.bin").read_bytes() == b"existing-face-index-runtime"
assert (root / "ocr" / "det.onnx").is_file()
manifest = json.loads((root / "manifest.json").read_text(encoding="utf-8"))
assert manifest["components"]["faces"]["version"] == "faces-existing"
assert manifest["components"]["ocr"]["version"] == "RapidOCR-PP-OCRv4"
assert "faces/sentinel.bin" in manifest["files"]
assert "ocr/det.onnx" in manifest["files"]
assert any("启用" in message for _, message in progress)
def test_component_activation_rolls_back_directory_and_manifest(tmp_path: Path, monkeypatch):
app, _ = _app(tmp_path)
manager = app.state.services.models
root = manager.settings.models_dir
old = root / "ocr"
old.mkdir(parents=True)
(old / "sentinel.bin").write_bytes(b"old-component")
original_manifest = {"version": "old", "components": {"ocr": {"version": "old-ocr"}}}
(root / "manifest.json").write_text(json.dumps(original_manifest), encoding="utf-8")
staged = tmp_path / "staged-ocr"
staged.mkdir()
for name in ("det.onnx", "rec.onnx", "cls.onnx"):
(staged / name).write_bytes(b"new")
monkeypatch.setattr(manager, "_write_manifest", lambda _manifest: (_ for _ in ()).throw(OSError("disk")))
with pytest.raises(OSError, match="disk"):
manager._activate_component("ocr", staged, {"version": "new-ocr"})
assert (old / "sentinel.bin").read_bytes() == b"old-component"
assert json.loads((root / "manifest.json").read_text(encoding="utf-8")) == original_manifest
def test_component_install_api_queues_missing_models_individually_and_reports_progress(tmp_path: Path):
app, headers = _app(tmp_path)
stale_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "stale"})
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post("/api/v1/models/install", headers=headers, json={"component": "all"})
assert response.status_code == 202
job_ids = response.json()["job_ids"]
assert response.json()["cancelled_audio_jobs"] == 1
# The deterministic hash visual backend is already ready in tests.
assert list(job_ids) == ["ocr", "faces", "audio"]
assert response.json()["job_id"] == job_ids["ocr"]
duplicate = await client.post("/api/v1/models/install", headers=headers, json={"component": "ocr"})
assert duplicate.json()["job_ids"]["ocr"] == job_ids["ocr"]
with app.state.services.db.transaction() as conn:
conn.execute(
"UPDATE jobs SET status='running',progress=.42,message='正在复制 OCR 权重' WHERE id=?",
(job_ids["ocr"],),
)
status = (await client.get("/api/v1/models", headers=headers)).json()
assert status["installations"]["ocr"] == {
"job_id": job_ids["ocr"],
"status": "running",
"progress": 0.42,
"message": "正在复制 OCR 权重",
"error": None,
}
custom = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "faces", "url": "https://example.invalid/models.tar.gz"},
)
assert custom.status_code == 400
with app.state.services.db.read() as conn:
audio = conn.execute("SELECT status FROM jobs WHERE id=?", (stale_audio,)).fetchone()
assert audio["status"] == "cancelled"
asyncio.run(scenario())
def test_audio_medium_install_is_independent_and_does_not_cancel_active_audio_work(tmp_path: Path):
app, headers = _app(tmp_path)
active_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "active"})
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "audio", "variant": "medium"},
)
assert response.status_code == 202
assert response.json()["cancelled_audio_jobs"] == 0
job_id = response.json()["job_ids"]["audio:medium"]
with app.state.services.db.read() as conn:
job = conn.execute("SELECT payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)).fetchone()
active = conn.execute("SELECT status FROM jobs WHERE id=?", (active_audio,)).fetchone()
assert json.loads(job["payload_json"]) == {"component": "audio", "variant": "medium"}
assert job["dedupe_key"] == "install-model:audio:medium"
assert active["status"] == "queued"
asyncio.run(scenario())
def test_pending_audio_install_cancels_stale_audio_jobs_on_startup(tmp_path: Path):
app, _ = _app(tmp_path)
stale_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "stale"})
app.state.services.jobs.enqueue("install_models", {"component": "audio"}, dedupe_key="install-model:audio")
assert app.state.services._cancel_audio_for_pending_model_install() == 1
with app.state.services.db.read() as conn:
row = conn.execute("SELECT status,message FROM jobs WHERE id=?", (stale_audio,)).fetchone()
assert row["status"] == "cancelled"
assert "模型安装优先" in row["message"]
def test_component_install_repairs_ready_runtime_without_redownloading_model(tmp_path: Path, monkeypatch):
app, headers = _app(tmp_path)
manager = app.state.services.models
monkeypatch.setattr(
manager,
"status",
lambda: {
"visual_ready": True,
"ocr_ready": False,
"faces_ready": False,
"audio_ready": False,
"operational_components": {"visual": False},
},
)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "visual"},
)
assert response.status_code == 202
job_id = response.json()["job_ids"]["visual"]
with app.state.services.db.read() as connection:
job = connection.execute(
"SELECT kind,payload_json,dedupe_key FROM jobs WHERE id=?",
(job_id,),
).fetchone()
assert job["kind"] == "prepare_ai_runtime"
assert json.loads(job["payload_json"]) == {"components": ["visual"]}
assert job["dedupe_key"] == "prepare-ai-runtime:visual"
asyncio.run(scenario())
def test_outdated_audio_repair_queues_real_model_reinstall(tmp_path: Path, monkeypatch):
app, headers = _app(tmp_path)
manager = app.state.services.models
monkeypatch.setattr(
manager,
"status",
lambda: {
"visual_ready": False,
"ocr_ready": False,
"faces_ready": False,
"audio_ready": True,
"operational_components": {"audio": False},
"component_health": {"audio": {"state": "outdated", "error": "需要转换为 FP16 模型"}},
},
)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "audio"},
)
assert response.status_code == 202
job_id = response.json()["job_ids"]["audio"]
with app.state.services.db.read() as connection:
job = connection.execute("SELECT kind,payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)).fetchone()
assert job["kind"] == "install_models"
assert json.loads(job["payload_json"]) == {"component": "audio"}
assert job["dedupe_key"] == "install-model:audio"
asyncio.run(scenario())
+320
View File
@@ -0,0 +1,320 @@
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")
+181
View File
@@ -0,0 +1,181 @@
from contextlib import nullcontext
from types import SimpleNamespace
import httpx
import pytest
from backend.imagefind.openlist_native import (
OpenListNativeService,
_remote_join,
_remote_parent,
normalize_task_state,
)
from backend.imagefind.remote import AlistClient
def test_remote_root_join_is_safe():
assert _remote_join(".", "", "/videos", "clip.mp4") == "videos/clip.mp4"
assert _remote_parent("clip.mp4") == ""
assert _remote_parent("videos/clip.mp4") == "videos"
assert normalize_task_state("succeeded") == 2
assert normalize_task_state("errored") == 5
assert normalize_task_state("7.0") == 7
def test_native_mapping_survives_library_only_deletion_by_falling_back_to_upload_history():
class Connection:
def execute(self, query, _parameters):
row = None
if "FROM uploads" in query:
row = {
"physical_path": "cloud/library/opaque.bin",
"physical_size_bytes": 356_793,
"size_bytes": 356_665,
}
return SimpleNamespace(fetchone=lambda: row)
service = object.__new__(OpenListNativeService)
service.db = SimpleNamespace(read=lambda: nullcontext(Connection()))
service.sources = SimpleNamespace(
get=lambda _source_id: {"config": {"storage_backend": "openlist_native"}}
)
assert service.physical_object_for_key("source", "ingest/movie.mp4") == (
"cloud/library/opaque.bin",
356_793,
356_665,
)
def test_alist_client_copy_and_application_errors():
calls: list[str] = []
def handler(request: httpx.Request) -> httpx.Response:
calls.append(request.url.path)
if request.url.path.endswith("/api/auth/login"):
return httpx.Response(200, json={"code": "200", "data": {"token": "test-token"}})
if request.url.path.endswith("/api/fs/get"):
body = request.content
if b"missing" in body:
return httpx.Response(200, json={"code": 404, "message": "object not found"})
return httpx.Response(200, json={"code": 200, "data": {"is_dir": False, "size": 4}})
if request.url.path.endswith("/api/fs/copy"):
return httpx.Response(200, json={"code": 200, "data": {"tasks": [{"id": "copy-1"}]}})
if request.url.path.endswith("/api/task/copy/info"):
return httpx.Response(200, json={"code": 200, "data": {"state": 2, "progress": 100}})
return httpx.Response(200, json={"code": 200, "data": {}})
client = AlistClient("http://openlist", "admin", "password")
client.client.close()
client.client = httpx.Client(transport=httpx.MockTransport(handler))
try:
assert client.object_info("video.mp4")["size"] == 4
assert client.object_info("missing") is None
assert client.copy_file("source/video.mp4", "stage/video.mp4") == ["copy-1"]
assert client.copy_task_info("copy-1")["state"] == 2
finally:
client.close()
assert calls.count("/api/auth/login") == 1
def test_native_catalog_uses_physical_object_but_keeps_logical_media_identity():
class Sources:
def get(self, source_id: str):
assert source_id == "source"
return {"config": {"mode": "direct"}}
def remote_access(self, source_id: str, key: str):
assert source_id == "source"
assert key == "合集/中文标题.mp4"
return "http://127.0.0.1:1234/%E5%90%88%E9%9B%86/video.mp4", "user", "secret", True
class Client:
def object_info(self, path: str):
assert path == "cloud/library/opaque-name.bin"
return {"is_dir": False, "size": 120, "modified": "2026-08-07T10:00:00Z", "sign": "etag"}
service = object.__new__(OpenListNativeService)
service.sources = Sources()
service.client = lambda _source_id: Client()
item = service.catalog_item(
"source",
"合集/中文标题.mp4",
{
"external_target_path": "cloud/library/opaque-name.bin",
"external_size_bytes": 120,
"size_bytes": 100,
"content_sha256": "a" * 64,
},
)
assert item.key == "合集/中文标题.mp4"
assert item.display_name == "中文标题.mp4"
assert item.size_bytes == 100
assert item.etag == "etag"
assert item.fingerprint
class _MovingClient:
def __init__(self, files: dict[str, int], *, fail_after_move: bool = False):
self.files = dict(files)
self.fail_after_move = fail_after_move
def object_info(self, path: str):
size = self.files.get(path)
return None if size is None else {"is_dir": False, "size": size}
def ensure_directory(self, _path: str):
return None
def move_file(self, source: str, target_directory: str, *, overwrite: bool = False):
target = f"{target_directory}/{source.rsplit('/', 1)[-1]}"
if not overwrite and target in self.files:
raise RuntimeError("target exists")
self.files[target] = self.files.pop(source)
if self.fail_after_move:
self.fail_after_move = False
raise TimeoutError("control-plane timeout")
def remove(self, path: str):
self.files.pop(path, None)
def _native_service(client: _MovingClient):
service = object.__new__(OpenListNativeService)
service.sources = SimpleNamespace(get=lambda _source_id: {})
service.configuration = lambda _source: SimpleNamespace(target_path="cloud/library", encrypted=False)
service.client = lambda _source_id, _source=None: client
return service
@pytest.mark.parametrize("ambiguous_timeout", [False, True])
def test_native_trash_and_restore_verify_provider_state(ambiguous_timeout: bool):
client = _MovingClient(
{"cloud/library/opaque.bin": 123},
fail_after_move=ambiguous_timeout,
)
service = _native_service(client)
trash_path = service.trash_object("source", "trash-id", "cloud/library/opaque.bin", 123)
assert trash_path == "cloud/library/.imagefind-native-trash/trash-id/opaque.bin"
assert client.files == {trash_path: 123}
client.fail_after_move = ambiguous_timeout
service.restore_object("source", trash_path, "cloud/library/opaque.bin", 123)
assert client.files == {"cloud/library/opaque.bin": 123}
def test_native_trash_restore_conflicts_are_non_destructive_and_purge_is_idempotent():
trash_path = "cloud/library/.imagefind-native-trash/trash-id/opaque.bin"
client = _MovingClient({"cloud/library/opaque.bin": 123, trash_path: 123})
service = _native_service(client)
with pytest.raises(RuntimeError, match="目标已存在"):
service.trash_object("source", "trash-id", "cloud/library/opaque.bin", 123)
with pytest.raises(FileExistsError, match="已被占用"):
service.restore_object("source", trash_path, "cloud/library/opaque.bin", 123)
assert client.files == {"cloud/library/opaque.bin": 123, trash_path: 123}
client.files.pop("cloud/library/opaque.bin")
service.purge_object("source", trash_path)
service.purge_object("source", trash_path)
assert client.files == {}
+804
View File
@@ -0,0 +1,804 @@
from __future__ import annotations
import asyncio
import contextlib
import json
import sys
import types
from pathlib import Path
import httpx
from imagefind import api as api_module
from imagefind.config import Settings
from imagefind.database import DatabaseTransientError, utcnow
from imagefind.main import create_app
def make_app(tmp_path: Path):
settings = Settings(data_dir=tmp_path / "data", embedding_backend="hash", upload_reserve_gb=0)
settings.prepare()
app = create_app(settings)
app.state.services.auth.setup("original administrator password")
_, token = app.state.services.auth.create_api_token("test")
return app, {"Authorization": f"Bearer {token}"}
def seed_videos(app) -> None:
now = utcnow()
with app.state.services.db.transaction() as conn:
conn.execute(
"INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)",
("source", "local", "Local", json.dumps({"path": "/media"}), now, now),
)
for index in range(3):
conn.execute(
"INSERT INTO videos(id,source_id,source_key,display_name,location,size_bytes,fingerprint,"
"duration_ms,status,available,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,'indexed',1,?,?)",
(
f"video-{index}",
"source",
f"video-{index}.mp4",
f"Video {index}",
f"/media/video-{index}.mp4",
100 + index,
f"fingerprint-{index}",
(index + 1) * 10_000,
now,
now,
),
)
conn.execute(
"INSERT INTO video_metadata(video_id,series,updated_at) VALUES('video-0','Series A',?),"
"('video-1','series a',?)",
(now, now),
)
conn.execute(
"INSERT INTO video_state(video_id,liked,favorited,progress_ms,completed,last_played_at,updated_at) "
"VALUES('video-0',1,0,3000,0,?,?),('video-1',0,1,10000,1,?,?)",
(now, now, now, now),
)
def test_video_list_batches_tag_items_in_one_read_connection(tmp_path: Path):
app, headers = make_app(tmp_path)
seed_videos(app)
service = app.state.services
now = utcnow()
with service.db.transaction() as conn:
conn.execute(
"INSERT INTO tag_groups(id,name,selection_mode,sort_order,created_at,updated_at) "
"VALUES('batch-group','类型','multi',1,?,?)",
(now, now),
)
conn.execute(
"INSERT INTO tags(id,group_id,name,created_at,updated_at) "
"VALUES('batch-tag','batch-group','测试标签',?,?)",
(now, now),
)
conn.executemany(
"INSERT INTO video_tags(video_id,tag_id) VALUES(?,'batch-tag')",
[(f"video-{index}",) for index in range(3)],
)
original_read = service.db.read
read_calls = []
@contextlib.contextmanager
def counted_read():
read_calls.append(1)
with original_read() as connection:
yield connection
service.db.read = counted_read
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.get("/api/v1/videos?limit=500", headers=headers)
assert response.status_code == 200
assert len(response.json()) == 3
assert all(item["tag_items"][0]["name"] == "测试标签" for item in response.json())
asyncio.run(scenario())
# The freshly issued API token is served by the bounded authentication
# cache, while one WAL read batches the complete video list/tag payload.
assert len(read_calls) == 1
def test_home_feed_groups_recent_collections_categories_and_unorganized(tmp_path: Path):
app, headers = make_app(tmp_path)
seed_videos(app)
now = utcnow()
with app.state.services.db.transaction() as conn:
conn.execute(
"INSERT INTO tag_groups(id,name,selection_mode,sort_order,created_at,updated_at) "
"VALUES('home-group','类型','multi',1,?,?)",
(now, now),
)
conn.execute(
"INSERT INTO tags(id,group_id,name,created_at,updated_at) VALUES"
"('home-tag-a','home-group','剧情',?,?),('home-tag-b','home-group','纪录',?,?),"
"('home-tag-small','home-group','单片',?,?)",
(now, now, now, now, now, now),
)
conn.executemany(
"INSERT INTO video_tags(video_id,tag_id) VALUES(?,?)",
[
("video-0", "home-tag-a"),
("video-1", "home-tag-a"),
("video-0", "home-tag-b"),
("video-1", "home-tag-b"),
("video-0", "home-tag-small"),
],
)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
assert (await client.get("/api/v1/home")).status_code == 401
collection = await client.post(
"/api/v1/collections", headers=headers, json={"name": "首页合集", "video_ids": ["video-0"]}
)
assert collection.status_code == 201
response = await client.get("/api/v1/home?item_limit=4&tag_limit=1", headers=headers)
assert response.status_code == 200
feed = response.json()
assert feed["video_count"] == 3
assert [item["id"] for item in feed["recent"]["items"]] == ["video-0", "video-1", "video-2"]
assert feed["collections"]["items"][0]["name"] == "首页合集"
assert len(feed["categories"]) == 1
assert feed["categories"][0]["title"] == "剧情"
assert feed["categories"][0]["total"] == 2
assert feed["unorganized"]["total"] == 1
assert [item["id"] for item in feed["unorganized"]["items"]] == ["video-2"]
with app.state.services.db.transaction() as conn:
conn.execute("DELETE FROM video_tags")
conn.execute("DELETE FROM video_metadata")
conn.execute("DELETE FROM collection_items")
conn.execute("DELETE FROM collection_videos")
conn.execute("DELETE FROM collections")
deduplicated = (await client.get("/api/v1/home", headers=headers)).json()
assert deduplicated["categories"] == []
assert deduplicated["collections"]["items"] == []
assert deduplicated["unorganized"] is None
asyncio.run(scenario())
def test_profile_validation_stats_filters_and_password_revocation(tmp_path: Path):
app, headers = make_app(tmp_path)
seed_videos(app)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
profile = await client.get("/api/v1/profile", headers=headers)
assert profile.status_code == 200
assert profile.json()["nickname"] == "管理员"
assert profile.json()["avatar"]["text"] == ""
assert profile.json()["avatar"]["accent"].startswith("#")
assert profile.json()["counts"] == {
"favorites": 1,
"likes": 1,
"history": 2,
"continue_watching": 1,
}
invalid = await client.patch("/api/v1/profile", headers=headers, json={"nickname": " "})
assert invalid.status_code == 400
changed = await client.patch("/api/v1/profile", headers=headers, json={"nickname": "NAS 管理员"})
assert changed.json()["avatar"]["text"] == "NA"
assert len((await client.get("/api/v1/videos?liked=true", headers=headers)).json()) == 1
assert len((await client.get("/api/v1/videos?favorite=true", headers=headers)).json()) == 1
assert len((await client.get("/api/v1/videos?continue_only=true", headers=headers)).json()) == 1
history = (await client.get("/api/v1/videos?played_only=true&sort=last_played", headers=headers)).json()
assert len(history) == 2
first = await client.post(
"/api/v1/auth/login", json={"password": "original administrator password"}
)
assert first.status_code == 200
old_cookie = first.cookies.get("imagefind_session")
wrong = await client.patch(
"/api/v1/profile/password",
headers=headers,
json={"current_password": "wrong", "new_password": "replacement administrator password"},
)
assert wrong.status_code == 400
changed_password = await client.patch(
"/api/v1/profile/password",
headers=headers,
json={
"current_password": "original administrator password",
"new_password": "replacement administrator password",
},
)
assert changed_password.status_code == 200
assert changed_password.json()["reauthenticate"] is True
client.cookies.set("imagefind_session", old_cookie)
revoked = await client.get("/api/v1/auth/me")
assert revoked.status_code == 401
assert (
await client.post("/api/v1/auth/login", json={"password": "replacement administrator password"})
).status_code == 200
asyncio.run(scenario())
def test_bulk_video_state_actions_are_atomic_and_preserve_unrelated_state(tmp_path: Path):
app, headers = make_app(tmp_path)
seed_videos(app)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
missing = await client.patch(
"/api/v1/videos/state/bulk",
headers=headers,
json={"video_ids": ["video-0", "missing"], "action": "unlike"},
)
assert missing.status_code == 404
assert len((await client.get("/api/v1/videos?liked=true", headers=headers)).json()) == 1
assert (
await client.patch(
"/api/v1/videos/state/bulk",
headers=headers,
json={"video_ids": ["video-0", "video-0"], "action": "unlike"},
)
).json() == {"updated": 1, "action": "unlike"}
await client.patch(
"/api/v1/videos/state/bulk",
headers=headers,
json={"video_ids": ["video-1"], "action": "unfavorite"},
)
await client.patch(
"/api/v1/videos/state/bulk",
headers=headers,
json={"video_ids": ["video-0", "video-1"], "action": "clear_history"},
)
assert (await client.get("/api/v1/videos?liked=true", headers=headers)).json() == []
assert (await client.get("/api/v1/videos?favorite=true", headers=headers)).json() == []
assert (await client.get("/api/v1/videos?played_only=true", headers=headers)).json() == []
asyncio.run(scenario())
def test_jobs_support_stable_numbered_pagination_and_legacy_limit(tmp_path: Path):
app, headers = make_app(tmp_path)
diagnostic_job = ""
for index in range(23):
diagnostic_job = app.state.services.jobs.enqueue("scan_source", {"index": index})
app.state.services.jobs.set_diagnostics(
diagnostic_job,
{
"requested_device": "GPU.0",
"actual_device": "CPU",
"fallback_scope": "job",
"fallback_reason": "low_quality_result",
"private_value": "must not be exposed",
},
)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
first = (await client.get("/api/v1/jobs?page=1&page_size=10", headers=headers)).json()
last = (await client.get("/api/v1/jobs?page=3&page_size=10", headers=headers)).json()
assert (first["page"], first["pages"], first["total"], len(first["items"])) == (1, 3, 23, 10)
assert (last["page"], len(last["items"])) == (3, 3)
assert not ({item["id"] for item in first["items"]} & {item["id"] for item in last["items"]})
diagnosed = next(item for item in first["items"] if item["id"] == diagnostic_job)
assert diagnosed["inference_diagnostics"] == {
"requested_device": "GPU.0",
"actual_device": "CPU",
"fallback_scope": "job",
"fallback_reason": "low_quality_result",
}
assert "diagnostics_json" not in diagnosed
legacy = (await client.get("/api/v1/jobs?limit=4", headers=headers)).json()
assert isinstance(legacy, list) and len(legacy) == 4
asyncio.run(scenario())
def test_job_diagnostics_are_best_effort_on_transient_database_conflict(
tmp_path: Path,
monkeypatch,
caplog,
):
app, _headers = make_app(tmp_path)
job_id = app.state.services.jobs.enqueue("scan_source", {"index": 1})
def conflict(*_args, **_kwargs):
raise DatabaseTransientError("temporary conflict")
monkeypatch.setattr(app.state.services.db, "write_with_retry", conflict)
app.state.services.jobs.set_diagnostics(job_id, {"actual_device": "GPU.0"})
assert "skipped diagnostics update" in caplog.text
def test_model_status_exposes_audio_circuit_breaker_state(tmp_path: Path):
app, headers = make_app(tmp_path)
accelerator = app.state.services.accelerator
accelerator.record_transient_failure("audio", "stall one", stage="inference_stall")
accelerator.record_transient_failure("audio", "stall two", stage="inference_stall")
accelerator.record_transient_failure("audio", "stall three", stage="inference_stall")
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.get("/api/v1/models", headers=headers)
assert response.status_code == 200
audio = response.json()["accelerator"]["components"]["audio"]
assert audio["circuit_state"] == "open"
assert audio["fallback_scope"] == "component"
assert audio["failure_count"] == 3
assert audio["retry_at"]
asyncio.run(scenario())
def test_upload_list_returns_last_snapshot_when_refresh_exceeds_budget(tmp_path: Path, monkeypatch):
app, headers = make_app(tmp_path)
cached = {"id": "cached-upload", "status": "receiving", "stage": "receiving"}
app.state.services.uploads._list_cache = [cached]
monkeypatch.setattr(api_module, "API_READ_DEADLINE_SECONDS", 0.01)
async def delayed(function, /, *args, **kwargs):
await asyncio.sleep(0.05)
return function(*args, **kwargs)
monkeypatch.setattr(api_module, "_background_api", delayed)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.get("/api/v1/uploads", headers=headers)
assert response.status_code == 200
assert response.json() == [cached]
# Let the shielded refresh finish so the test loop closes cleanly.
await asyncio.sleep(0.06)
asyncio.run(scenario())
def test_upload_page_returns_matching_snapshot_when_refresh_exceeds_budget(tmp_path: Path, monkeypatch):
app, headers = make_app(tmp_path)
cached = {
"items": [{"id": "cached-page-upload", "status": "completed", "stage": "done"}],
"status_items": [{"id": "cached-live-upload", "status": "receiving", "stage": "receiving"}],
"active_count": 1,
"failed_count": 0,
"page": 2,
"page_size": 10,
"total": 11,
"pages": 2,
}
app.state.services.uploads._page_cache[(2, 10)] = cached
monkeypatch.setattr(api_module, "API_READ_DEADLINE_SECONDS", 0.01)
async def delayed(function, /, *args, **kwargs):
await asyncio.sleep(0.05)
return function(*args, **kwargs)
monkeypatch.setattr(api_module, "_background_api", delayed)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.get("/api/v1/uploads?page=2&page_size=10", headers=headers)
assert response.status_code == 200
assert response.json() == cached
await asyncio.sleep(0.06)
asyncio.run(scenario())
def test_failed_jobs_can_be_retried_without_overwriting_history(tmp_path: Path):
app, headers = make_app(tmp_path)
failed_id = app.state.services.jobs.enqueue(
"scan_source",
{"source_id": "source-retry"},
dedupe_key="scan:source-retry",
)
with app.state.services.db.transaction() as connection:
connection.execute(
"UPDATE jobs SET status='failed',progress=.35,error='temporary failure',finished_at=? WHERE id=?",
(utcnow(), failed_id),
)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(f"/api/v1/jobs/{failed_id}/retry", headers=headers, json={})
assert response.status_code == 202
retried_id = response.json()["job_id"]
assert response.json()["retried_from"] == failed_id
duplicate = await client.post(f"/api/v1/jobs/{failed_id}/retry", headers=headers, json={})
assert duplicate.json()["job_id"] == retried_id
active = await client.post(f"/api/v1/jobs/{retried_id}/retry", headers=headers, json={})
assert active.status_code == 409
missing = await client.post("/api/v1/jobs/missing/retry", headers=headers, json={})
assert missing.status_code == 404
with app.state.services.db.read() as connection:
failed = connection.execute("SELECT status,error FROM jobs WHERE id=?", (failed_id,)).fetchone()
retried = connection.execute(
"SELECT kind,payload_json,dedupe_key,status,attempts FROM jobs WHERE id=?",
(retried_id,),
).fetchone()
assert dict(failed) == {"status": "failed", "error": "temporary failure"}
assert retried["kind"] == "scan_source"
assert json.loads(retried["payload_json"]) == {"source_id": "source-retry"}
assert retried["dedupe_key"] == "scan:source-retry"
assert (retried["status"], retried["attempts"]) == ("queued", 0)
asyncio.run(scenario())
def test_failed_jobs_can_be_retried_in_bulk_by_lane_without_replaying_superseded_history(
tmp_path: Path,
):
app, headers = make_app(tmp_path)
jobs = app.state.services.jobs
scan_failed = jobs.enqueue(
"scan_source", {"source_id": "source-bulk"}, dedupe_key="scan:source-bulk"
)
audio_failed = jobs.enqueue(
"transcribe_audio", {"video_id": "video-bulk"}, dedupe_key="audio:video-bulk"
)
superseded_failed = jobs.enqueue(
"index_video", {"video_id": "video-superseded"}, dedupe_key="index:video-superseded"
)
with app.state.services.db.transaction() as connection:
connection.execute(
"UPDATE jobs SET status='failed',error='temporary failure',finished_at=? "
"WHERE id IN (?,?,?)",
(utcnow(), scan_failed, audio_failed, superseded_failed),
)
superseding_job = jobs.enqueue(
"index_video", {"video_id": "video-superseded"}, dedupe_key="index:video-superseded"
)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
scan_page = await client.get(
"/api/v1/jobs?page=1&page_size=10&lane=scan", headers=headers
)
assert scan_page.status_code == 200
assert scan_page.json()["retryable_failed_count"] == 1
retried_scan = await client.post(
"/api/v1/jobs/retry-failed?lane=scan", headers=headers, json={}
)
assert retried_scan.status_code == 202
assert retried_scan.json()["retried"] == 1
assert retried_scan.json()["lane"] == "scan"
duplicate = await client.post(
"/api/v1/jobs/retry-failed?lane=scan", headers=headers, json={}
)
assert duplicate.status_code == 202
assert duplicate.json()["retried"] == 0
all_page = await client.get("/api/v1/jobs?page=1&page_size=10", headers=headers)
assert all_page.json()["retryable_failed_count"] == 1
retried_audio = await client.post("/api/v1/jobs/retry-failed", headers=headers, json={})
assert retried_audio.status_code == 202
assert retried_audio.json()["retried"] == 1
with app.state.services.db.read() as connection:
old_statuses = {
row["id"]: row["status"]
for row in connection.execute(
"SELECT id,status FROM jobs WHERE id IN (?,?,?)",
(scan_failed, audio_failed, superseded_failed),
).fetchall()
}
active = connection.execute(
"SELECT id,kind,dedupe_key FROM jobs WHERE status='queued' ORDER BY created_at"
).fetchall()
assert old_statuses == {
scan_failed: "failed",
audio_failed: "failed",
superseded_failed: "failed",
}
assert [row["dedupe_key"] for row in active].count("index:video-superseded") == 1
assert {row["kind"] for row in active} >= {
"scan_source",
"transcribe_audio",
"index_video",
}
assert superseding_job in {row["id"] for row in active}
asyncio.run(scenario())
def test_series_aggregate_merge_bulk_and_filter(tmp_path: Path):
app, headers = make_app(tmp_path)
seed_videos(app)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
rows = (await client.get("/api/v1/series", headers=headers)).json()
assert rows[0]["video_count"] == 2
assert rows[0]["duration_ms"] == 30_000
renamed = await client.patch(
"/api/v1/series", headers=headers, json={"name": "SERIES A", "new_name": "Merged"}
)
assert renamed.json()["updated"] == 2
assigned = await client.post(
"/api/v1/videos/series/bulk",
headers=headers,
json={"video_ids": ["video-2"], "series": "Merged"},
)
assert assigned.json()["updated"] == 1
assert len((await client.get("/api/v1/videos?series=Merged", headers=headers)).json()) == 3
removed = await client.post(
"/api/v1/videos/series/bulk",
headers=headers,
json={"video_ids": ["video-2"], "series": None},
)
assert removed.json()["series"] is None
asyncio.run(scenario())
def test_collections_support_single_membership_order_cover_and_dissolve(tmp_path: Path):
app, headers = make_app(tmp_path)
seed_videos(app)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
first = await client.post(
"/api/v1/collections",
headers=headers,
json={"name": "第一季", "description": "按顺序播放", "video_ids": ["video-0", "video-1"]},
)
assert first.status_code == 201
first_id = first.json()["id"]
second = await client.post(
"/api/v1/collections",
headers=headers,
json={"name": "临时合集", "video_ids": ["video-2"]},
)
second_id = second.json()["id"]
moved = await client.post(
f"/api/v1/collections/{first_id}/videos",
headers=headers,
json={"video_ids": ["video-2"]},
)
assert moved.json()["updated"] == 1
assert (await client.get(f"/api/v1/collections/{second_id}", headers=headers)).json()[
"video_count"
] == 0
detail = (await client.get(f"/api/v1/collections/{first_id}", headers=headers)).json()
assert [video["id"] for video in detail["videos"]] == ["video-0", "video-1", "video-2"]
reordered = await client.patch(
f"/api/v1/collections/{first_id}/videos/order",
headers=headers,
json={"video_ids": ["video-2", "video-0", "video-1"]},
)
assert reordered.json()["updated"] == 3
invalid = await client.patch(
f"/api/v1/collections/{first_id}/videos/order",
headers=headers,
json={"video_ids": ["video-0"]},
)
assert invalid.status_code == 400
changed = await client.patch(
f"/api/v1/collections/{first_id}",
headers=headers,
json={"description": "更新后的简介", "cover_video_id": "video-1"},
)
assert changed.json()["cover_video_id"] == "video-1"
filtered = (
await client.get(
f"/api/v1/videos?collection_id={first_id}&sort=collection",
headers=headers,
)
).json()
assert [video["id"] for video in filtered] == ["video-2", "video-0", "video-1"]
assert all(video["collection_id"] == first_id for video in filtered)
dissolved = await client.delete(f"/api/v1/collections/{first_id}", headers=headers)
assert dissolved.json()["detached_videos"] == 3
assert len((await client.get("/api/v1/videos?limit=20", headers=headers)).json()) == 3
with app.state.services.db.read() as conn:
assert conn.execute("SELECT count(*) FROM collection_videos").fetchone()[0] == 0
asyncio.run(scenario())
def test_proxy_is_encrypted_redacted_and_excluded_from_model_status(tmp_path: Path):
app, headers = make_app(tmp_path)
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.patch(
"/api/v1/models/config",
headers=headers,
json={
"hf_endpoint": "https://hf-mirror.example",
"proxy_url": "http://proxy.example:8080",
"proxy_username": "proxy-user",
"proxy_password": "proxy-secret",
},
)
assert response.status_code == 200
assert response.json()["proxy"] == {
"enabled": True,
"url": "http://proxy.example:8080",
"username": "proxy-user",
"has_password": True,
}
status = (await client.get("/api/v1/models", headers=headers)).json()
assert "proxy-secret" not in json.dumps(status)
assert "password" not in status["proxy"]
assert status["accelerator"]["policy"] == "gpu_preferred"
assert status["accelerator"]["cpu_threads"] == 1
assert set(status["accelerator"]["components"]) == {"visual", "ocr", "faces", "audio"}
assert status["pip"]["index_url"] == "https://pypi.tuna.tsinghua.edu.cn/simple"
disabled = await client.patch(
"/api/v1/models/config",
headers=headers,
json={"proxy_enabled": False},
)
assert disabled.json()["proxy"]["enabled"] is False
assert disabled.json()["proxy"]["has_password"] is True
assert app.state.services.models._httpx_proxy() is None
assert app.state.services.models._hf_proxies() is None
sources = await client.patch(
"/api/v1/models/config",
headers=headers,
json={
"pip_index_url": "https://pypi.example/simple/",
"pytorch_index_url": "https://torch.example/cpu/",
},
)
assert sources.json()["pip"] == {
"index_url": "https://pypi.example/simple",
"pytorch_index_url": "https://torch.example/cpu",
}
rejected = await client.patch(
"/api/v1/models/config",
headers=headers,
json={"pip_index_url": "https://user:secret@pypi.example/simple"},
)
assert rejected.status_code == 400
asyncio.run(scenario())
stored = app.state.services.db.setting("model_proxy_config")
assert "proxy-secret" not in json.dumps(stored)
assert app.state.services.models.proxy_config(include_password=True)["password"] == "proxy-secret"
def test_huggingface_downloads_ignore_inherited_proxy_when_switch_is_off_for_non_mirror_endpoint(
tmp_path: Path, monkeypatch
):
app, _ = make_app(tmp_path)
manager = app.state.services.models
factories = []
class Session:
def __init__(self):
self.trust_env = True
self.proxies = {}
monkeypatch.setitem(sys.modules, "requests", types.SimpleNamespace(Session=Session))
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
types.SimpleNamespace(
configure_http_backend=lambda *, backend_factory: factories.append(backend_factory)
),
)
manager.settings.model_hf_endpoint = "https://huggingface.co"
manager.set_proxy_config("http://proxy.example:8080", enabled=False)
manager._configure_hf_http_backend()
direct = factories[-1]()
assert direct.trust_env is False
assert direct.proxies == {}
manager.set_proxy_config("http://proxy.example:8080", enabled=True)
manager._configure_hf_http_backend()
proxied = factories[-1]()
assert proxied.trust_env is False
assert proxied.proxies == {
"http": "http://proxy.example:8080",
"https": "http://proxy.example:8080",
}
def test_hf_mirror_downloads_stay_direct_even_when_model_proxy_is_enabled(tmp_path: Path, monkeypatch):
app, _ = make_app(tmp_path)
manager = app.state.services.models
factories = []
class Session:
def __init__(self):
self.trust_env = True
self.proxies = {}
monkeypatch.setitem(sys.modules, "requests", types.SimpleNamespace(Session=Session))
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
types.SimpleNamespace(configure_http_backend=lambda *, backend_factory: factories.append(backend_factory)),
)
manager.set_proxy_config("http://proxy.example:8080", enabled=True)
manager.settings.model_hf_endpoint = "https://hf-mirror.com"
assert manager._hf_proxies() is None
manager._configure_hf_http_backend()
direct = factories[-1]()
assert direct.trust_env is False
assert direct.proxies == {}
def test_model_uninstall_preserves_indexes_and_reports_usage(tmp_path: Path):
app, headers = make_app(tmp_path)
seed_videos(app)
root = app.state.services.settings.models_dir
(root / "ocr").mkdir(parents=True)
for name in ("det.onnx", "rec.onnx", "cls.onnx"):
(root / "ocr" / name).write_bytes(b"model-contents")
(root / "manifest.json").write_text(
json.dumps({"version": "v1", "components": {"ocr": {"version": "ocr-v1"}}}), encoding="utf-8"
)
with app.state.services.db.transaction() as conn:
conn.execute(
"INSERT INTO frames(id,video_id,timestamp_ms,segment_start_ms,segment_end_ms,thumbnail_path,"
"vector_blob,created_at) VALUES('frame','video-0',0,0,1,'thumb',?,?)",
(b"existing-vector", utcnow()),
)
app.state.services.accelerator.mark_ready("ocr", "CPU")
assert app.state.services.accelerator.status()["components"]["ocr"]["state"] == "ready"
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
deleted = await client.delete("/api/v1/models/ocr", headers=headers)
assert deleted.status_code == 200
assert deleted.json()["freed_bytes"] > 0
usage = (await client.get("/api/v1/storage/usage?refresh=true", headers=headers)).json()
assert set(usage["categories"]) == {
"database",
"ai_index",
"models",
"thumbnails",
"preview_cache",
"remote_cache",
"upload_staging",
"other",
}
assert usage["categories"]["ai_index"]["mode"] == "pgvector"
asyncio.run(scenario())
assert not (root / "ocr").exists()
assert app.state.services.accelerator.status()["components"]["ocr"]["state"] == "not_loaded"
with app.state.services.db.read() as conn:
assert conn.execute("SELECT vector_blob FROM frames WHERE id='frame'").fetchone()[0] == b"existing-vector"
def test_model_uninstall_conflicts_with_install_job(tmp_path: Path):
app, headers = make_app(tmp_path)
app.state.services.jobs.enqueue("install_models", {}, dedupe_key="install-models")
async def scenario() -> None:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.delete("/api/v1/models/all", headers=headers)
assert response.status_code == 409
asyncio.run(scenario())
File diff suppressed because it is too large Load Diff
+881
View File
@@ -0,0 +1,881 @@
from __future__ import annotations
import hashlib
import io
import json
import subprocess
import sys
import tarfile
import types
import zipfile
from pathlib import Path
import imagefind.runtime as runtime_module
import pytest
from imagefind.config import Settings
from imagefind.runtime import AIDependencyManager, RuntimeToolManager
def _settings(tmp_path: Path, requirements: Path | None = None) -> Settings:
if requirements is not None:
requirements.mkdir(parents=True, exist_ok=True)
(requirements / "constraints-cp312.txt").write_text("# test constraints\n")
for component in ("visual", "ocr", "faces", "audio"):
(requirements / f"{component}.txt").touch(exist_ok=True)
settings = Settings(data_dir=tmp_path / "data", runtime_requirements_dir=requirements)
settings.prepare()
return settings
def test_ai_dependency_failure_preserves_previous_layer_and_redacts_proxy_password(tmp_path: Path, monkeypatch):
requirements = tmp_path / "requirements"
requirements.mkdir()
(requirements / "audio.txt").write_text("torch==test\n")
settings = _settings(tmp_path, requirements)
current = settings.ai_site_path
current.mkdir()
(current / "sentinel").write_text("previous")
(current / ".imagefind-runtime.json").write_text(
json.dumps({"lock_version": "old", "components": ["visual"]})
)
manager = AIDependencyManager(
settings,
lambda **_: {
"enabled": True,
"url": "http://proxy.invalid:8080",
"username": "imagefind",
"password": "top/secret?",
},
)
def failed_install(*_args, **_kwargs):
return subprocess.CompletedProcess([], 1, "", "proxy authentication failed: top%2Fsecret%3F")
monkeypatch.setattr(runtime_module.subprocess, "run", failed_install)
with pytest.raises(RuntimeError) as error:
manager.ensure("audio")
assert "top/secret?" not in str(error.value)
assert "top%2Fsecret%3F" not in str(error.value)
assert "***" in str(error.value)
assert (current / "sentinel").read_text() == "previous"
assert manager.status()["visual"]["state"] == "missing"
assert manager.status()["audio"]["state"] == "error"
assert "***" in manager.status()["audio"]["error"]
def test_ai_runtime_uses_configured_pip_sources_and_obeys_proxy_switch(tmp_path: Path, monkeypatch):
requirements = tmp_path / "requirements"
requirements.mkdir()
(requirements / "visual.txt").write_text("torch==test\n")
settings = _settings(tmp_path, requirements)
settings.pip_index_url = "https://pypi.tuna.example/simple"
settings.pytorch_index_url = "https://torch.example/cpu"
proxy = {
"enabled": False,
"url": "http://proxy.example:8080",
"username": "",
"password": "",
}
manager = AIDependencyManager(settings, lambda **_: proxy)
environments: list[dict[str, str]] = []
def successful_install(command, **kwargs):
environments.append(kwargs["env"])
return subprocess.CompletedProcess(command, 0, "", "")
monkeypatch.setenv("HTTPS_PROXY", "http://inherited-proxy.invalid:9999")
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
manager.ensure("visual")
assert environments[0]["PIP_INDEX_URL"] == "https://pypi.tuna.example/simple"
assert environments[0]["PIP_EXTRA_INDEX_URL"] == "https://torch.example/cpu"
assert "HTTPS_PROXY" not in environments[0]
assert environments[0]["PIP_DEFAULT_TIMEOUT"] == "60"
assert environments[0]["PIP_RETRIES"] == "3"
def test_ai_runtime_injects_enabled_proxy_without_logging_password(tmp_path: Path, monkeypatch):
requirements = tmp_path / "requirements"
requirements.mkdir()
(requirements / "audio.txt").write_text("torch==test\n")
settings = _settings(tmp_path, requirements)
manager = AIDependencyManager(
settings,
lambda **_: {
"enabled": True,
"url": "http://proxy.example:8080",
"username": "runtime-user",
"password": "runtime/password",
},
)
environments: list[dict[str, str]] = []
def successful_install(command, **kwargs):
environments.append(kwargs["env"])
return subprocess.CompletedProcess(command, 0, "", "")
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
manager.ensure("audio")
assert environments[0]["HTTPS_PROXY"] == (
"http://runtime-user:runtime%2Fpassword@proxy.example:8080"
)
def test_ai_runtime_stamp_binds_python_abi_and_lock_digests(tmp_path: Path, monkeypatch):
requirements = tmp_path / "requirements"
settings = _settings(tmp_path, requirements)
(requirements / "visual.txt").write_text("openvino==test\n")
manager = AIDependencyManager(settings)
monkeypatch.setattr(
runtime_module.subprocess,
"run",
lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "", ""),
)
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
manager.ensure("visual")
stamp = json.loads(manager.stamp_path.read_text())
record = stamp["components"]["visual"]
assert stamp["schema"] == runtime_module.AI_RUNTIME_SCHEMA
assert record["python_abi"] == manager._python_abi()
assert len(record["requirements_sha256"]) == 64
assert len(record["constraints_sha256"]) == 64
assert manager.status()["visual"]["state"] == "ready"
(requirements / "constraints-cp312.txt").write_text("# comment-only release change\n")
assert manager.status()["visual"]["state"] == "ready"
(requirements / "constraints-cp312.txt").write_text("openvino==changed\n")
assert manager.status()["visual"]["state"] == "missing"
def test_schema_two_component_stamps_from_0319_all_remain_ready(tmp_path: Path, monkeypatch):
requirements = tmp_path / "requirements"
settings = _settings(tmp_path, requirements)
project_requirements = Path(__file__).resolve().parents[1] / "requirements" / "runtime-ai"
current_constraints = (project_requirements / "constraints-cp312.txt").read_text()
_, remainder = current_constraints.split("\n", 1)
(requirements / "constraints-cp312.txt").write_text(
"# ImageFind 0.3.19 AI runtime lock for fnOS Python 3.12.\n" + remainder
)
for component in runtime_module.AI_IMPORTS:
(requirements / f"{component}.txt").write_text(
(project_requirements / f"{component}.txt").read_text()
)
manager = AIDependencyManager(settings)
constraints_digest = runtime_module._sha256(requirements / "constraints-cp312.txt")
records = {}
for component in runtime_module.AI_IMPORTS:
requirement_digest = runtime_module._sha256(requirements / f"{component}.txt")
combined = hashlib.sha256()
combined.update(manager._python_abi().encode())
combined.update(requirement_digest.encode())
combined.update(constraints_digest.encode())
for value in runtime_module.AI_NO_DEPENDENCIES.get(component, ()):
combined.update(value.encode())
records[component] = {
"python_abi": manager._python_abi(),
"requirements_sha256": requirement_digest,
"constraints_sha256": constraints_digest,
"digest": combined.hexdigest(),
"runtime_version": "0.3.19",
}
manager.current.mkdir(parents=True, exist_ok=True)
manager.stamp_path.write_text(
json.dumps(
{
"schema": 2,
"runtime_version": "0.3.19",
"components": records,
}
)
)
monkeypatch.setattr(
runtime_module.subprocess,
"run",
lambda *_args, **_kwargs: pytest.fail("compatible runtime must not invoke pip"),
)
assert {name: value["state"] for name, value in manager.status().items()} == {
component: "ready" for component in runtime_module.AI_IMPORTS
}
for component in runtime_module.AI_IMPORTS:
manager.ensure(component)
def test_vad_constraint_upgrade_only_invalidates_audio_runtime(tmp_path: Path):
requirements = tmp_path / "requirements"
settings = _settings(tmp_path, requirements)
project_requirements = Path(__file__).resolve().parents[1] / "requirements" / "runtime-ai"
constraints = runtime_module._requirements_lines(project_requirements / "constraints-cp312.txt")
old_constraints = tuple(line for line in constraints if line != "webrtcvad-wheels==2.0.14")
(requirements / "constraints-cp312.txt").write_text("\n".join(constraints) + "\n")
for component in runtime_module.AI_IMPORTS:
(requirements / f"{component}.txt").write_text(
(project_requirements / f"{component}.txt").read_text()
)
manager = AIDependencyManager(settings)
prior_constraints_digest = hashlib.sha256(("\n".join(old_constraints) + "\n").encode()).hexdigest()
records = {}
for component in runtime_module.AI_IMPORTS:
component_lines = runtime_module._requirements_lines(requirements / f"{component}.txt")
if component == "audio":
component_lines = tuple(line for line in component_lines if line != "webrtcvad-wheels==2.0.14")
requirement_digest = hashlib.sha256(("\n".join(component_lines) + "\n").encode()).hexdigest()
combined = hashlib.sha256()
combined.update(manager._python_abi().encode())
combined.update(requirement_digest.encode())
combined.update(prior_constraints_digest.encode())
for value in runtime_module.AI_NO_DEPENDENCIES.get(component, ()):
combined.update(value.encode())
records[component] = {
"digest_format": runtime_module.AI_RUNTIME_DIGEST_FORMAT,
"python_abi": manager._python_abi(),
"requirements_sha256": requirement_digest,
"constraints_sha256": prior_constraints_digest,
"digest": combined.hexdigest(),
}
manager.current.mkdir(parents=True, exist_ok=True)
manager.stamp_path.write_text(json.dumps({"schema": 3, "components": records}))
states = {name: value["state"] for name, value in manager.status().items()}
assert states == {"visual": "ready", "ocr": "ready", "faces": "ready", "audio": "missing"}
def test_runtime_tools_prefer_system_binary(tmp_path: Path, monkeypatch):
settings = _settings(tmp_path)
manager = RuntimeToolManager(settings)
monkeypatch.setattr(runtime_module.shutil, "which", lambda name: "/system/bin/ffmpeg" if name == "ffmpeg" else None)
monkeypatch.setattr(manager, "_works", lambda path, name: path == "/system/bin/ffmpeg" and name == "ffmpeg")
monkeypatch.setattr(manager, "_install_release", lambda _name: pytest.fail("private fallback should not download"))
assert manager.ffmpeg() == "/system/bin/ffmpeg"
def test_bundled_ai_runtime_is_ready_and_never_invokes_pip(tmp_path: Path, monkeypatch):
settings = _settings(tmp_path)
settings.bundled_ai_runtime = True
manager = AIDependencyManager(settings)
imported: list[str] = []
monkeypatch.setattr(runtime_module.importlib, "import_module", lambda name: imported.append(name))
monkeypatch.setattr(runtime_module.subprocess, "run", lambda *_args, **_kwargs: pytest.fail("pip must not run"))
manager.ensure("audio")
assert imported == list(runtime_module.AI_IMPORTS["audio"])
assert manager.status()["audio"]["source"] == "bundled"
assert manager.environment_status()["components"] == list(runtime_module.AI_IMPORTS)
def test_runtime_tools_use_bundled_fallback_without_downloading(tmp_path: Path, monkeypatch):
settings = _settings(tmp_path)
settings.bundled_tools_dir = tmp_path / "app-bin"
settings.bundled_tools_dir.mkdir()
bundled = settings.bundled_tools_dir / "ffmpeg"
bundled.write_bytes(b"bundled")
manager = RuntimeToolManager(settings)
monkeypatch.setattr(runtime_module.shutil, "which", lambda _name: None)
monkeypatch.setattr(manager, "_works", lambda path, name: path == str(bundled) and name == "ffmpeg")
monkeypatch.setattr(manager, "_install_release", lambda _name: pytest.fail("fallback should not download"))
assert manager.ffmpeg() == str(bundled)
def test_ocr_runtime_uses_headless_opencv_without_desktop_dependency(tmp_path: Path, monkeypatch):
requirements = tmp_path / "requirements"
requirements.mkdir()
(requirements / "ocr.txt").write_text("opencv-python-headless==4.13.0.92\n")
settings = _settings(tmp_path, requirements)
manager = AIDependencyManager(settings)
commands: list[list[str]] = []
def successful_install(command, **_kwargs):
commands.append(command)
return subprocess.CompletedProcess(command, 0, "", "")
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
manager.ensure("ocr")
assert "--requirement" in commands[0]
assert "opencv-python-headless==4.13.0.92" in (requirements / "ocr.txt").read_text()
assert "opencv-python==" not in (requirements / "ocr.txt").read_text()
assert "--no-deps" in commands[1]
assert commands[1][-1] == "rapidocr-onnxruntime==1.4.4"
def test_installing_one_missing_component_preserves_and_reuses_compatible_runtime(
tmp_path: Path, monkeypatch
):
requirements = tmp_path / "requirements"
settings = _settings(tmp_path, requirements)
project_requirements = Path(__file__).resolve().parents[1] / "requirements" / "runtime-ai"
for component in runtime_module.AI_IMPORTS:
(requirements / f"{component}.txt").write_text(
(project_requirements / f"{component}.txt").read_text()
)
manager = AIDependencyManager(settings)
commands: list[list[str]] = []
def successful_install(command, **_kwargs):
commands.append(command)
return subprocess.CompletedProcess(command, 0, "", "")
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
manager.ensure("visual")
(manager.current / "preserved-package").write_text("keep")
assert manager.status()["visual"]["state"] == "ready"
assert manager.status()["audio"]["state"] == "missing"
assert manager.status()["ocr"]["state"] == "missing"
assert manager.status()["faces"]["state"] == "missing"
commands.clear()
manager.ensure("ocr")
assert len(commands) == 2
assert commands[0][-1] == str(requirements / "ocr.txt")
assert commands[1][-1] == "rapidocr-onnxruntime==1.4.4"
assert all(str(requirements / "visual.txt") not in command for command in commands)
assert (manager.current / "preserved-package").read_text() == "keep"
assert {name: value["state"] for name, value in manager.status().items()} == {
"visual": "ready",
"ocr": "ready",
"faces": "ready",
"audio": "missing",
}
monkeypatch.setattr(
runtime_module.subprocess,
"run",
lambda *_args, **_kwargs: pytest.fail("covered components must not invoke pip"),
)
manager.ensure("faces")
commands.clear()
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
manager.ensure("audio")
assert len(commands) == 1
assert commands[0][-1] == str(requirements / "audio.txt")
def test_preserved_models_trigger_runtime_repair_and_accelerator_refresh(tmp_path: Path, monkeypatch):
from imagefind.models import ModelManager
requirements = tmp_path / "requirements"
requirements.mkdir()
settings = _settings(tmp_path, requirements)
class Accelerator:
def __init__(self):
self.refreshes = 0
def refresh(self):
self.refreshes += 1
accelerator = Accelerator()
embeddings = types.SimpleNamespace(accelerator=accelerator)
manager = ModelManager(
settings,
embeddings,
types.SimpleNamespace(),
types.SimpleNamespace(),
)
manager.component_versions = lambda: {
"visual": "visual-v1",
"ocr": "ocr-v1",
"faces": None,
"audio": None,
}
manager.runtime_dependencies.status = lambda: {
"visual": {"state": "missing"},
"ocr": {"state": "ready"},
"faces": {"state": "missing"},
"audio": {"state": "missing"},
}
assert accelerator.refreshes == 1
assert manager.missing_runtime_components() == ["visual"]
ensured = []
monkeypatch.setattr(
manager.runtime_dependencies,
"ensure",
lambda component, progress=None: ensured.append((component, progress)),
)
manager.ensure_runtime("visual")
assert ensured == [("visual", None)]
assert accelerator.refreshes == 2
def test_runtime_repair_resets_stale_component_fallback(tmp_path: Path, monkeypatch):
from imagefind.models import ModelManager
requirements = tmp_path / "requirements"
requirements.mkdir()
settings = _settings(tmp_path, requirements)
class Accelerator:
def __init__(self):
self.refreshes = 0
def refresh(self):
self.refreshes += 1
resets: list[str] = []
accelerator = Accelerator()
embeddings = types.SimpleNamespace(accelerator=accelerator)
ocr = types.SimpleNamespace(reset=lambda: resets.append("ocr"))
manager = ModelManager(settings, embeddings, ocr, types.SimpleNamespace())
monkeypatch.setattr(manager.runtime_dependencies, "ensure", lambda *_args, **_kwargs: None)
manager.ensure_runtime("ocr")
assert resets == ["ocr"]
assert accelerator.refreshes == 2
def test_unavailable_accelerator_makes_installed_component_repairable(tmp_path: Path):
from imagefind.models import ModelManager
settings = _settings(tmp_path)
accelerator = types.SimpleNamespace(
refresh=lambda: None,
status=lambda: {"components": {"visual": {"state": "unavailable"}}},
)
manager = ModelManager(
settings,
types.SimpleNamespace(accelerator=accelerator),
types.SimpleNamespace(),
types.SimpleNamespace(),
)
versions = {"visual": "visual-v1", "ocr": None, "faces": None, "audio": None}
runtime = {name: {"state": "ready"} for name in versions}
health = {name: {"state": "ready", "error": None} for name in versions}
assert manager.operational_components(versions, runtime, health)["visual"] is False
def test_visual_image_export_uses_clip_openvino_task_instead_of_generic_text_backend(
tmp_path: Path, monkeypatch
):
from imagefind.models import ModelManager
settings = _settings(tmp_path)
source = tmp_path / "clip-image"
source.mkdir()
(source / "modules.json").write_text(
'[{"idx":0,"path":"0_CLIPModel","type":"sentence_transformers.models.CLIPModel"}]'
)
transformer_source = source / "0_CLIPModel"
transformer_source.mkdir()
(transformer_source / "config.json").write_text("{}")
destination = tmp_path / "exported-image"
calls = {}
class ExportConfig:
pass
def main_export(**kwargs):
calls["export"] = kwargs
path = kwargs["output"]
path.mkdir(parents=True)
(path / "openvino_model.xml").write_text("<xml/>")
(path / "openvino_model.bin").write_bytes(b"model")
optimum = types.ModuleType("optimum")
optimum_exporters = types.ModuleType("optimum.exporters")
optimum_exporters_openvino = types.ModuleType("optimum.exporters.openvino")
optimum_exporters_openvino.main_export = main_export
optimum_intel = types.ModuleType("optimum.intel")
optimum_openvino = types.ModuleType("optimum.intel.openvino")
optimum_configuration = types.ModuleType("optimum.intel.openvino.configuration")
optimum_configuration.OVConfig = ExportConfig
for name, module in {
"optimum": optimum,
"optimum.exporters": optimum_exporters,
"optimum.exporters.openvino": optimum_exporters_openvino,
"optimum.intel": optimum_intel,
"optimum.intel.openvino": optimum_openvino,
"optimum.intel.openvino.configuration": optimum_configuration,
}.items():
monkeypatch.setitem(sys.modules, name, module)
class Accelerator:
@staticmethod
def refresh():
pass
@staticmethod
def ov_config(device):
assert device == "CPU"
return {"INFERENCE_NUM_THREADS": 2}
embeddings = types.SimpleNamespace(accelerator=Accelerator())
manager = ModelManager(settings, embeddings, types.SimpleNamespace(), types.SimpleNamespace())
manager._export_visual_image_openvino(source, destination)
assert calls["export"] == {
"model_name_or_path": str(transformer_source.resolve()),
"output": destination / "openvino",
"task": "zero-shot-image-classification",
"library_name": "transformers",
"local_files_only": True,
"ov_config": calls["export"]["ov_config"],
}
assert isinstance(calls["export"]["ov_config"], ExportConfig)
assert (destination / "modules.json").is_file()
assert (destination / "0_CLIPModel" / "config.json").is_file()
assert (destination / "openvino" / "openvino_model.xml").is_file()
def test_visual_text_export_uses_transformer_module_and_explicit_library(
tmp_path: Path, monkeypatch
):
from imagefind.models import ModelManager
settings = _settings(tmp_path)
source = tmp_path / "clip-text"
source.mkdir()
(source / "modules.json").write_text(
'[{"idx":0,"path":"0_Transformer","type":"sentence_transformers.models.Transformer"}]'
)
transformer_source = source / "0_Transformer"
transformer_source.mkdir()
(transformer_source / "config.json").write_text("{}")
destination = tmp_path / "exported-text"
calls = {}
class ExportConfig:
pass
def main_export(**kwargs):
calls["export"] = kwargs
path = kwargs["output"]
path.mkdir(parents=True)
(path / "openvino_model.xml").write_text("<xml/>")
(path / "openvino_model.bin").write_bytes(b"model")
optimum = types.ModuleType("optimum")
optimum_exporters = types.ModuleType("optimum.exporters")
optimum_exporters_openvino = types.ModuleType("optimum.exporters.openvino")
optimum_exporters_openvino.main_export = main_export
optimum_intel = types.ModuleType("optimum.intel")
optimum_openvino = types.ModuleType("optimum.intel.openvino")
optimum_configuration = types.ModuleType("optimum.intel.openvino.configuration")
optimum_configuration.OVConfig = ExportConfig
for name, module in {
"optimum": optimum,
"optimum.exporters": optimum_exporters,
"optimum.exporters.openvino": optimum_exporters_openvino,
"optimum.intel": optimum_intel,
"optimum.intel.openvino": optimum_openvino,
"optimum.intel.openvino.configuration": optimum_configuration,
}.items():
monkeypatch.setitem(sys.modules, name, module)
embeddings = types.SimpleNamespace(
accelerator=types.SimpleNamespace(refresh=lambda: None)
)
manager = ModelManager(settings, embeddings, types.SimpleNamespace(), types.SimpleNamespace())
manager._export_visual_text_openvino(source, destination)
assert calls["export"] == {
"model_name_or_path": str(transformer_source.resolve()),
"output": destination / "0_Transformer" / "openvino",
"task": "feature-extraction",
"library_name": "transformers",
"local_files_only": True,
"ov_config": calls["export"]["ov_config"],
}
assert isinstance(calls["export"]["ov_config"], ExportConfig)
assert (destination / "modules.json").is_file()
assert (destination / "0_Transformer" / "config.json").is_file()
assert (destination / "0_Transformer" / "openvino" / "openvino_model.xml").is_file()
def test_ai_and_model_sizes_are_cached_for_status_polling(tmp_path: Path, monkeypatch):
from imagefind.models import ModelManager
requirements = tmp_path / "requirements"
settings = _settings(tmp_path, requirements)
runtime = AIDependencyManager(settings)
runtime_calls = []
monkeypatch.setattr(
runtime_module,
"_directory_size",
lambda path: runtime_calls.append(path) or 123,
)
assert runtime._current_size() == 123
assert runtime._current_size() == 123
assert runtime_calls == [runtime.current]
embeddings = types.SimpleNamespace(
accelerator=types.SimpleNamespace(refresh=lambda: None)
)
manager = ModelManager(settings, embeddings, types.SimpleNamespace(), types.SimpleNamespace())
model_calls = []
monkeypatch.setattr(
manager,
"_directory_size",
lambda path: model_calls.append(path) or 456,
)
assert manager._component_sizes() == {
"visual": 456,
"ocr": 456,
"faces": 456,
"audio": 456,
}
assert manager._component_sizes()["visual"] == 456
assert len(model_calls) == 4
def test_preserved_model_files_queue_runtime_repair_after_upgrade(tmp_path: Path):
from imagefind.main import create_app
requirements = tmp_path / "requirements"
requirements.mkdir()
settings = _settings(tmp_path, requirements)
(settings.models_dir / "visual" / "image").mkdir(parents=True)
(settings.models_dir / "visual" / "text").mkdir(parents=True)
(settings.models_dir / "manifest.json").write_text(
json.dumps({"components": {"visual": {"version": "visual-v1"}}})
)
services = create_app(settings).state.services
assert services.models.missing_runtime_components() == ["visual"]
job_id = services.queue_missing_ai_runtime()
with services.db.read() as connection:
job = connection.execute(
"SELECT kind,payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)
).fetchone()
assert job["kind"] == "prepare_ai_runtime"
assert json.loads(job["payload_json"]) == {"components": ["visual"]}
assert job["dedupe_key"] == "prepare-ai-runtime:visual"
def test_private_ffmpeg_and_rclone_releases_coexist_and_failed_update_is_atomic(tmp_path: Path, monkeypatch):
settings = _settings(tmp_path)
manager = RuntimeToolManager(settings)
ffmpeg = b"ffmpeg-test"
ffprobe = b"ffprobe-test"
rclone = b"rclone-test"
ffmpeg_archive = tmp_path / "ffmpeg.tar.xz"
with tarfile.open(ffmpeg_archive, "w:xz") as archive:
for name, content in (("release/ffmpeg", ffmpeg), ("release/ffprobe", ffprobe)):
member = tarfile.TarInfo(name)
member.size = len(content)
archive.addfile(member, io.BytesIO(content))
rclone_archive = tmp_path / "rclone.zip"
with zipfile.ZipFile(rclone_archive, "w") as archive:
archive.writestr("release/rclone", rclone)
releases = {
"ffmpeg": {
"version": "test",
"url": str(ffmpeg_archive),
"files": {
"ffmpeg": hashlib.sha256(ffmpeg).hexdigest(),
"ffprobe": hashlib.sha256(ffprobe).hexdigest(),
},
},
"rclone": {
"version": "test",
"url": str(rclone_archive),
"files": {"rclone": hashlib.sha256(rclone).hexdigest()},
},
}
monkeypatch.setattr(runtime_module, "TOOL_RELEASES", releases)
def local_download(url: str, destination: Path):
destination.write_bytes(Path(url).read_bytes())
monkeypatch.setattr(manager, "_download", local_download)
manager._install_release("ffmpeg")
manager._install_release("rclone")
assert (manager.root / "ffmpeg").read_bytes() == ffmpeg
assert (manager.root / "ffprobe").read_bytes() == ffprobe
assert (manager.root / "rclone").read_bytes() == rclone
releases["rclone"]["files"]["rclone"] = "0" * 64
with pytest.raises(RuntimeError, match="SHA-256"):
manager._install_release("rclone")
assert (manager.root / "ffmpeg").read_bytes() == ffmpeg
assert (manager.root / "rclone").read_bytes() == rclone
def test_audio_export_declares_transformers_library_and_fp16_configuration(tmp_path: Path, monkeypatch):
calls: dict[str, object] = {}
class ExportConfig:
def __init__(self, *, dtype):
self.dtype = dtype
class Processor:
@classmethod
def from_pretrained(cls, source, *, local_files_only):
calls["processor_source"] = source
calls["processor_local"] = local_files_only
return cls()
def save_pretrained(self, destination):
calls["processor_destination"] = destination
_write_minimal_audio_processor_files(destination)
def main_export(**kwargs):
calls["export"] = kwargs
_write_minimal_audio_export(kwargs["output"])
modules = {
"optimum": types.ModuleType("optimum"),
"optimum.exporters": types.ModuleType("optimum.exporters"),
"optimum.exporters.openvino": types.ModuleType("optimum.exporters.openvino"),
"optimum.intel": types.ModuleType("optimum.intel"),
"optimum.intel.openvino": types.ModuleType("optimum.intel.openvino"),
"optimum.intel.openvino.configuration": types.ModuleType("optimum.intel.openvino.configuration"),
"transformers": types.ModuleType("transformers"),
}
modules["optimum.exporters.openvino"].main_export = main_export
modules["optimum.intel.openvino.configuration"].OVConfig = ExportConfig
modules["transformers"].AutoProcessor = Processor
for name, module in modules.items():
monkeypatch.setitem(sys.modules, name, module)
from imagefind.models import ModelManager
source = tmp_path / "source"
destination = tmp_path / "output"
ModelManager._export_audio_model(source, destination)
export = calls["export"]
assert export["library_name"] == "transformers"
assert export["task"] == "automatic-speech-recognition-with-past"
assert export["local_files_only"] is True
assert export["ov_config"].dtype == "fp16"
assert calls["processor_destination"] == destination
def _write_minimal_audio_export(root: Path, *, include_cache: bool = False, marker: str = "main") -> None:
root.mkdir(parents=True, exist_ok=True)
for name in (
"openvino_encoder_model.xml",
"openvino_decoder_model.xml",
):
(root / name).write_text(f"<{marker}/>")
for name in (
"openvino_encoder_model.bin",
"openvino_decoder_model.bin",
):
(root / name).write_bytes(marker.encode())
if include_cache:
(root / "openvino_decoder_with_past_model.xml").write_text(f"<{marker}-cache/>")
(root / "openvino_decoder_with_past_model.bin").write_bytes(marker.encode())
def _write_minimal_audio_processor_files(root: Path) -> None:
root.mkdir(parents=True, exist_ok=True)
for name in ("config.json", "preprocessor_config.json", "tokenizer_config.json", "tokenizer.json"):
(root / name).write_text("{}", encoding="utf-8")
def _install_audio_export_modules(monkeypatch, main_export, *, fallback_model=None):
class ExportConfig:
def __init__(self, *, dtype):
self.dtype = dtype
class Processor:
@classmethod
def from_pretrained(cls, source, *, local_files_only):
assert local_files_only is True
return cls()
def save_pretrained(self, destination):
_write_minimal_audio_processor_files(destination)
modules = {
"optimum": types.ModuleType("optimum"),
"optimum.exporters": types.ModuleType("optimum.exporters"),
"optimum.exporters.openvino": types.ModuleType("optimum.exporters.openvino"),
"optimum.intel": types.ModuleType("optimum.intel"),
"optimum.intel.openvino": types.ModuleType("optimum.intel.openvino"),
"optimum.intel.openvino.configuration": types.ModuleType("optimum.intel.openvino.configuration"),
"transformers": types.ModuleType("transformers"),
}
modules["optimum.exporters.openvino"].main_export = main_export
modules["optimum.intel.openvino.configuration"].OVConfig = ExportConfig
if fallback_model is not None:
modules["optimum.intel.openvino"].OVModelForSpeechSeq2Seq = fallback_model
modules["transformers"].AutoProcessor = Processor
for name, module in modules.items():
monkeypatch.setitem(sys.modules, name, module)
def test_audio_export_keeps_main_export_when_optional_cache_export_fails(
tmp_path: Path, monkeypatch
):
def main_export(**kwargs):
_write_minimal_audio_export(kwargs["output"], marker="main")
class FailingFallback:
@classmethod
def from_pretrained(cls, *_args, **_kwargs):
raise ValueError("cache export unavailable")
_install_audio_export_modules(monkeypatch, main_export, fallback_model=FailingFallback)
from imagefind.models import ModelManager
destination = tmp_path / "audio"
ModelManager._export_audio_model(tmp_path / "source", destination)
assert (destination / "openvino_encoder_model.xml").read_text() == "<main/>"
assert not (destination / "openvino_decoder_with_past_model.xml").exists()
assert (destination / "tokenizer_config.json").is_file()
def test_audio_export_replaces_main_export_only_when_cache_export_is_complete(
tmp_path: Path, monkeypatch
):
def main_export(**kwargs):
_write_minimal_audio_export(kwargs["output"], marker="main")
class SuccessfulFallback:
@classmethod
def from_pretrained(cls, *_args, **_kwargs):
return cls()
def save_pretrained(self, destination):
_write_minimal_audio_export(destination, include_cache=True, marker="cache")
_install_audio_export_modules(monkeypatch, main_export, fallback_model=SuccessfulFallback)
from imagefind.models import ModelManager
destination = tmp_path / "audio"
ModelManager._export_audio_model(tmp_path / "source", destination)
assert (destination / "openvino_encoder_model.xml").read_text() == "<cache/>"
assert (destination / "openvino_decoder_with_past_model.xml").read_text() == "<cache-cache/>"
assert (destination / "tokenizer_config.json").is_file()
def test_audio_export_main_export_failure_reports_root_cause(tmp_path: Path, monkeypatch):
def main_export(**_kwargs):
raise ValueError("library name could not be inferred")
_install_audio_export_modules(monkeypatch, main_export)
from imagefind.models import ModelManager
with pytest.raises(RuntimeError, match="音频模型转换阶段失败:ValueError.*library name"):
ModelManager._export_audio_model(tmp_path / "source", tmp_path / "audio")
+73
View File
@@ -0,0 +1,73 @@
from pathlib import Path
import httpx
from imagefind.sources import LocalConnector, WebDavConnector
def test_local_connector_ignores_symlinks_and_non_video(tmp_path: Path):
(tmp_path / "a.mp4").write_bytes(b"video")
(tmp_path / "note.txt").write_text("no")
outside = tmp_path.parent / "outside.mkv"
outside.write_bytes(b"outside")
(tmp_path / "link.mkv").symlink_to(outside)
items = list(LocalConnector(str(tmp_path)).items())
assert [item.key for item in items] == ["a.mp4"]
def test_webdav_depth_one_and_external_href_rejection():
xml = b"""<?xml version="1.0"?>
<d:multistatus xmlns:d="DAV:">
<d:response>
<d:href>/videos/</d:href><d:propstat><d:prop>
<d:resourcetype><d:collection/></d:resourcetype>
</d:prop><d:status>HTTP/1.1 200 OK</d:status></d:propstat>
</d:response>
<d:response>
<d:href>/videos/movie.mp4</d:href><d:propstat><d:prop>
<d:resourcetype/><d:getcontentlength>123</d:getcontentlength><d:getetag>abc</d:getetag>
</d:prop><d:status>HTTP/1.1 200 OK</d:status></d:propstat>
</d:response>
<d:response>
<d:href>https://evil.example/leak.mp4</d:href><d:propstat><d:prop>
<d:resourcetype/>
</d:prop><d:status>HTTP/1.1 200 OK</d:status></d:propstat>
</d:response>
</d:multistatus>"""
def handler(request: httpx.Request):
assert request.method == "PROPFIND"
assert request.headers["Depth"] == "1"
return httpx.Response(207, content=xml)
connector = WebDavConnector("https://dav.example/videos/", "user", "pass")
connector.client.close()
connector.client = httpx.Client(transport=httpx.MockTransport(handler))
items = list(connector.items())
connector.close()
assert len(items) == 1
assert items[0].key == "movie.mp4"
assert items[0].etag == "abc"
def test_webdav_tolerates_alist_xml_with_extra_payload():
xml = b"""alist notice
<D:multistatus xmlns:D="DAV:">
<D:response>
<D:href>/dav/movie.mp4</D:href><D:propstat><D:prop>
<D:resourcetype/><D:getcontentlength>42</D:getcontentlength>
</D:prop><D:status>HTTP/1.1 200 OK</D:status></D:propstat>
</D:response>
</D:multistatus><ignored/>"""
def handler(_request: httpx.Request):
return httpx.Response(207, content=xml)
connector = WebDavConnector("https://dav.example/dav/", "user", "pass")
connector.client.close()
connector.client = httpx.Client(transport=httpx.MockTransport(handler))
items = list(connector.items())
connector.close()
assert len(items) == 1
assert items[0].key == "movie.mp4"
assert items[0].size_bytes == 42
+120
View File
@@ -0,0 +1,120 @@
from __future__ import annotations
from types import SimpleNamespace
from imagefind.speech_quality import (
normalize_language,
pcm16_voiced_regions,
select_language,
transcript_quality,
)
def test_language_policy_treats_stream_metadata_as_weak_hint() -> None:
assert select_language(["en", "zh"], policy="zh_priority") == "zh"
assert select_language(["en", "zh"], policy="auto") == "en"
assert select_language(["en", "en"], policy="zh_priority", stream_language="zh-CN") == "zh"
assert select_language([], policy="zh_priority") == "zh"
assert select_language([], policy="zh_priority", stream_language="eng") == "zh"
assert select_language([], policy="auto", stream_language="eng") == "en"
assert normalize_language("<|zh|>") == "zh"
def test_chinese_priority_rejects_weak_or_implausible_language_detection() -> None:
assert select_language(["jw", "nn", "jw"], policy="zh_priority") == "zh"
assert select_language(["en"], policy="zh_priority") == "zh"
assert select_language(["en", "en", "zh"], policy="zh_priority") == "zh"
def test_chinese_priority_does_not_trust_short_korean_or_japanese_detection() -> None:
# Whisper's one-token detector commonly confuses short/noisy Mandarin
# clips with ko/ja. Choosing that language makes the whole decode wrong.
assert select_language(["ko", "ko", "ko"], policy="zh_priority") == "zh"
assert select_language(["ja", "ja"], policy="zh_priority") == "zh"
assert select_language(["jw", "jw"], policy="auto") == "jw"
def test_quality_accepts_normal_chinese_english_and_mixed_transcripts() -> None:
samples = (
("这是一个正常的中文语音识别结果,包含完整的句子和清晰的信息。", "zh"),
("This is a clear English transcript with enough useful information.", "en"),
("今天我们 discuss ImageFind 的 audio search 功能。", "zh"),
)
for text, language in samples:
quality = transcript_quality(text, 12_000, speech_ratio=0.7, expected_language=language)
assert quality.accepted is True
assert quality.score >= 0.9
assert quality.flags == ()
def test_quality_rejects_repeated_wrong_script_and_replacement_characters() -> None:
repeated = transcript_quality(
"කකකකකකකකකකකකකකකකකකකක",
12_000,
speech_ratio=0.7,
expected_language="zh",
)
damaged = transcript_quality(
"这是一段损坏的字幕内容",
8_000,
speech_ratio=0.6,
expected_language="zh",
)
assert repeated.accepted is False
assert {"repeated_characters", "unexpected_script"}.issubset(repeated.flags)
assert damaged.accepted is False
assert "invalid_characters" in damaged.flags
def test_quality_rejects_full_english_hallucination_when_chinese_is_forced() -> None:
result = transcript_quality(
"This entire segment was decoded in the wrong language.",
8_000,
speech_ratio=0.8,
expected_language="zh",
)
assert result.accepted is False
assert "language_script_conflict" in result.flags
def test_quality_rejects_short_phrase_hallucination_but_accepts_silence() -> None:
hallucination = transcript_quality(
"谢谢观看",
30_000,
speech_ratio=0.01,
expected_language="zh",
)
silence = transcript_quality("", 30_000, speech_ratio=0.0, expected_language="zh")
assert hallucination.accepted is False
assert "short_hallucination" in hallucination.flags
assert silence.accepted is True
assert silence.units == 0
def test_quality_does_not_reject_sparse_legitimate_speech_by_chunk_duration() -> None:
quality = transcript_quality("今天天气不错", 30_000, speech_ratio=0.08, expected_language="zh")
assert quality.accepted is True
assert "too_little_text" not in quality.flags
def test_vad_regions_bridge_short_gaps_and_skip_short_noise(monkeypatch) -> None:
decisions = [False] * 4 + [True] * 10 + [False] * 8 + [True] * 10 + [False] * 20 + [True] * 3
class Detector:
def __init__(self, _mode):
self.index = 0
def is_speech(self, _frame, _sample_rate):
value = decisions[self.index]
self.index += 1
return value
monkeypatch.setitem(__import__("sys").modules, "webrtcvad", SimpleNamespace(Vad=Detector))
raw = b"\0\0" * 480 * len(decisions)
regions = pcm16_voiced_regions(raw, 16_000)
assert len(regions) == 1
start, end, ratio = regions[0]
assert start < 4 * 480
assert end > 32 * 480
assert 0 < ratio < 1
+166
View File
@@ -0,0 +1,166 @@
import asyncio
import json
from pathlib import Path
import httpx
from imagefind.config import Settings
from imagefind.database import utcnow
from imagefind.main import create_app
from imagefind.vectors import pack_vector
def _seed_video(app, tmp_path: Path) -> str:
now = utcnow()
video_id = "tag-video"
with app.state.services.db.transaction() as conn:
conn.execute(
"INSERT INTO sources(id,kind,name,config_json,created_at,updated_at) VALUES(?,?,?,?,?,?)",
("tag-source", "local", "标签测试库", json.dumps({"path": str(tmp_path)}), now, now),
)
conn.execute(
"INSERT INTO videos(id,source_id,source_key,display_name,location,fingerprint,duration_ms,width,status,"
"available,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)",
(
video_id,
"tag-source",
"ABC-123.mp4",
"ABC-123.mp4",
str(tmp_path / "ABC-123.mp4"),
"tag-fingerprint",
100_000,
3840,
"indexed",
1,
now,
now,
),
)
conn.execute(
"INSERT INTO frames(id,video_id,timestamp_ms,segment_start_ms,segment_end_ms,thumbnail_path,"
"vector_blob,created_at) VALUES(?,?,?,?,?,?,?,?)",
("tag-frame", video_id, 20_000, 18_000, 24_000, str(tmp_path / "frame.webp"), pack_vector([1, 0]), now),
)
return video_id
def test_tag_groups_suggestions_state_and_preferences(tmp_path: Path):
app = create_app(Settings(data_dir=tmp_path, embedding_backend="hash", scan_interval_seconds=86400))
service = app.state.services
video_id = _seed_video(app, tmp_path)
service.vectors.upsert("frames", "tag-frame", [1, 0], {"video_id": video_id})
_, token = service.auth.create_api_token("test")
headers = {"Authorization": f"Bearer {token}"}
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
group = await client.post(
"/api/v1/tag-groups",
headers=headers,
json={"name": "类型", "selection_mode": "single"},
)
assert group.status_code == 201
group_id = group.json()["id"]
tag_ids = []
for name in ("剧情", "访谈"):
response = await client.post(
"/api/v1/tags",
headers=headers,
json={"group_id": group_id, "name": name},
)
assert response.status_code == 201
tag_ids.append(response.json()["id"])
conflict = await client.patch(
f"/api/v1/videos/{video_id}/metadata",
headers=headers,
json={"tag_ids": tag_ids},
)
assert conflict.status_code == 400
saved = await client.patch(
f"/api/v1/videos/{video_id}/metadata",
headers=headers,
json={"tag_ids": [tag_ids[0]]},
)
assert saved.status_code == 200
visual_group = await client.post(
"/api/v1/tag-groups",
headers=headers,
json={"name": "场景", "selection_mode": "multi"},
)
visual_tag = await client.post(
"/api/v1/tags",
headers=headers,
json={
"group_id": visual_group.json()["id"],
"name": "室内",
"ai_enabled": True,
"ai_method": "visual",
"ai_description": "室内房间",
"ai_threshold": 0.5,
},
)
service.embeddings.encode_text = lambda _text: [1, 0]
service.tag_suggestions.analyze("test-job", [video_id], force=True)
suggestions = await client.get(
f"/api/v1/videos/{video_id}/tag-suggestions",
headers=headers,
)
assert suggestions.status_code == 200
assert suggestions.json()[0]["tag_id"] == visual_tag.json()["id"]
assert suggestions.json()[0]["evidence"][0]["timestamp_ms"] == 20_000
decision = await client.post(
"/api/v1/tag-suggestions/decide",
headers=headers,
json={"suggestion_ids": [suggestions.json()[0]["id"]], "action": "accept"},
)
assert decision.status_code == 200
state = await client.patch(
f"/api/v1/videos/{video_id}/state",
headers=headers,
json={"liked": True, "favorited": True, "progress_ms": 42_000},
)
assert state.status_code == 200
assert state.json()["progress_ms"] == 42_000
videos = await client.get("/api/v1/videos?favorite=true&min_width=3800", headers=headers)
assert len(videos.json()) == 1
assert videos.json()[0]["tag_items"]
assert videos.json()[0]["liked"] == 1
defaults = await client.get("/api/v1/preferences", headers=headers)
assert defaults.status_code == 200
assert defaults.json()["home_video_columns"] == 2
assert defaults.json()["theme"] == "system"
preferences = await client.patch(
"/api/v1/preferences",
headers=headers,
json={
"autoplay": False,
"theme": "dark",
"home_video_columns": 1,
"upload_paths": {"tag-source": "待整理/2026"},
},
)
assert preferences.status_code == 200
assert preferences.json()["autoplay"] is False
assert preferences.json()["theme"] == "dark"
assert preferences.json()["home_video_columns"] == 1
assert preferences.json()["upload_paths"]["tag-source"] == "待整理/2026"
invalid_columns = await client.patch(
"/api/v1/preferences",
headers=headers,
json={"home_video_columns": 3},
)
assert invalid_columns.status_code == 422
invalid_theme = await client.patch(
"/api/v1/preferences",
headers=headers,
json={"theme": "midnight"},
)
assert invalid_theme.status_code == 422
asyncio.run(scenario())
+30
View File
@@ -0,0 +1,30 @@
from imagefind.subtitles import parse_subtitles, timestamp_ms
from imagefind.text import fts_query, normalize_text, search_tokens
def test_chinese_and_english_tokens_are_stable():
tokens = search_tokens("红色汽车 Red-Car 订单A12")
assert "红色" in tokens
assert "汽车" in tokens
assert "red-car" in tokens
assert "a12" in tokens
assert normalize_text(" ABC ") == "abc"
assert '"红色"' in fts_query("红色")
def test_srt_and_vtt_parser():
cues = parse_subtitles(
"""WEBVTT
00:00:01.200 --> 00:00:03.400
<b>你好</b> world
00:01:10,000 --> 00:01:11,500
第二行
"""
)
assert [(cue.start_ms, cue.end_ms, cue.text) for cue in cues] == [
(1200, 3400, "你好 world"),
(70000, 71500, "第二行"),
]
assert timestamp_ms("1:02:03.045") == 3_723_045
@@ -0,0 +1,858 @@
from __future__ import annotations
import asyncio
import base64
import hashlib
import sys
import time
from pathlib import Path
import httpx
import pytest
from imagefind import webdav as webdav_module
from imagefind.collections import assign_videos, set_collection_tags
from imagefind.config import Settings
from imagefind.database import utcnow
from imagefind.main import create_app
from imagefind.webdav import _path_lock
@pytest.fixture(autouse=True)
def _inline_webdav_background_io(monkeypatch):
"""Avoid Python 3.13's sandbox-only to_thread selector deadlock.
fnOS runs Python 3.12 and production keeps these calls in worker threads;
the test suite still exercises the same write, hash and state logic inline.
"""
async def inline(_app, function, /, *args, **kwargs):
return function(*args, **kwargs)
monkeypatch.setattr(webdav_module, "_background_io", inline)
def _app(tmp_path: Path, **overrides):
settings = Settings(
data_dir=tmp_path / "data",
embedding_backend="hash",
upload_chunk_mb=1,
upload_staging_gb=1,
upload_reserve_gb=0,
**overrides,
)
settings.prepare()
app = create_app(settings)
media = tmp_path / "media"
media.mkdir()
source_id = app.state.services.sources.add_local("媒体库", str(media))
app.state.services.storage.set_writable(source_id, True)
_, token = app.state.services.auth.create_api_token("test")
return app, media, source_id, token
def _seed_taxonomy(app, source_id: str) -> tuple[str, str, str, str]:
now = utcnow()
with app.state.services.db.transaction() as conn:
group_id = "kind"
conn.execute(
"INSERT INTO tag_groups(id,name,selection_mode,sort_order,created_at,updated_at) "
"VALUES(?,?,'single',0,?,?)",
(group_id, "类型", now, now),
)
conn.executemany(
"INSERT INTO tags(id,group_id,name,created_at,updated_at) VALUES(?,?,?,?,?)",
(("movie", group_id, "电影", now, now), ("course", group_id, "课程", now, now)),
)
collection_id = "collection"
conn.execute(
"INSERT INTO collections(id,name,description,created_at,updated_at) VALUES(?,?, '',?,?)",
(collection_id, "旅行", now, now),
)
video_id = "existing-video"
conn.execute(
"INSERT INTO videos(id,source_id,source_key,display_name,location,size_bytes,fingerprint,"
"duration_ms,status,created_at,updated_at) VALUES(?,?,?,?,?,1,?,60000,'ready',?,?)",
(video_id, source_id, "existing.mp4", "existing.mp4", "", "fingerprint", now, now),
)
return collection_id, video_id, "movie", "course"
def test_collection_default_tags_apply_now_and_on_future_membership(tmp_path: Path):
app, _, source_id, _ = _app(tmp_path)
collection_id, existing_id, movie_tag, course_tag = _seed_taxonomy(app, source_id)
now = utcnow()
with app.state.services.db.transaction() as conn:
assign_videos(conn, collection_id, [existing_id])
conn.execute("INSERT INTO video_tags(video_id,tag_id) VALUES(?,?)", (existing_id, course_tag))
applied = set_collection_tags(conn, collection_id, [movie_tag])
assert applied == 1
future_id = "future-video"
conn.execute(
"INSERT INTO videos(id,source_id,source_key,display_name,location,size_bytes,fingerprint,"
"duration_ms,status,created_at,updated_at) VALUES(?,?,?,?,?,1,?,60000,'ready',?,?)",
(future_id, source_id, "future.mp4", "future.mp4", "", "future", now, now),
)
assign_videos(conn, collection_id, [future_id])
assert conn.execute(
"SELECT tag_id FROM video_tags WHERE video_id=?", (existing_id,)
).fetchone()[0] == movie_tag
assert conn.execute(
"SELECT tag_id FROM video_tags WHERE video_id=?", (future_id,)
).fetchone()[0] == movie_tag
set_collection_tags(conn, collection_id, [])
assign_videos(conn, None, [future_id])
assert conn.execute(
"SELECT tag_id FROM video_tags WHERE video_id=?", (future_id,)
).fetchone()[0] == movie_tag
def test_upload_metadata_is_applied_when_catalogued(tmp_path: Path):
app, media, source_id, _ = _app(tmp_path)
collection_id, _, movie_tag, _ = _seed_taxonomy(app, source_id)
payload = b"video" * 100
upload = app.state.services.uploads.create(
source_id,
"imports",
"trip.mp4",
len(payload),
collection_id=collection_id,
tag_ids=[movie_tag],
)
app.state.services.uploads.receive_chunk(
upload["id"], 0, payload, hashlib.sha256(payload).hexdigest()
)
app.state.services.uploads.complete(upload["id"])
queued = app.state.services.uploads._get(upload["id"])
app.state.services.uploads.transfer(queued["job_id"], upload["id"])
assert (media / "imports" / "trip.mp4").read_bytes() == payload
app.state.services.scanner.refresh_path("refresh", source_id, "imports/trip.mp4", upload["id"])
with app.state.services.db.read() as conn:
video = conn.execute(
"SELECT id FROM videos WHERE source_id=? AND source_key='imports/trip.mp4'", (source_id,)
).fetchone()
assert conn.execute(
"SELECT collection_id FROM collection_videos WHERE video_id=?", (video["id"],)
).fetchone()[0] == collection_id
assert conn.execute(
"SELECT tag_id FROM video_tags WHERE video_id=?", (video["id"],)
).fetchone()[0] == movie_tag
def test_video_marker_crud_and_duration_validation(tmp_path: Path):
app, _, source_id, token = _app(tmp_path)
_, video_id, _, _ = _seed_taxonomy(app, source_id)
headers = {"Authorization": f"Bearer {token}"}
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
created = await client.post(
f"/api/v1/videos/{video_id}/markers",
headers=headers,
json={"position_ms": 12345},
)
assert created.status_code == 201
marker = created.json()
updated = await client.patch(
f"/api/v1/videos/{video_id}/markers/{marker['id']}",
headers=headers,
json={"title": "精彩片段"},
)
assert updated.json()["title"] == "精彩片段"
assert (await client.get(f"/api/v1/videos/{video_id}/markers", headers=headers)).json()[0][
"position_ms"
] == 12345
invalid = await client.post(
f"/api/v1/videos/{video_id}/markers",
headers=headers,
json={"position_ms": 60001},
)
assert invalid.status_code == 400
assert (
await client.delete(
f"/api/v1/videos/{video_id}/markers/{marker['id']}", headers=headers
)
).status_code == 204
asyncio.run(scenario())
def test_collection_tree_supports_arbitrary_depth_move_and_group_promotion(tmp_path: Path):
app, _, source_id, token = _app(tmp_path)
collection_id, video_id, _, _ = _seed_taxonomy(app, source_id)
headers = {"Authorization": f"Bearer {token}"}
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
chapter = (
await client.post(
f"/api/v1/collections/{collection_id}/groups",
headers=headers,
json={"name": "第一章"},
)
).json()
section = (
await client.post(
f"/api/v1/collections/{collection_id}/groups",
headers=headers,
json={"name": "第一节", "parent_id": chapter["id"]},
)
).json()
assigned = await client.post(
f"/api/v1/collections/{collection_id}/videos",
headers=headers,
json={"video_ids": [video_id], "parent_id": section["id"]},
)
assert assigned.status_code == 200
detail = (await client.get(f"/api/v1/collections/{collection_id}", headers=headers)).json()
assert detail["items"][0]["children"][0]["children"][0]["video_id"] == video_id
assert detail["videos"][0]["collection_path"] == ["第一章", "第一节"]
cycle = await client.patch(
f"/api/v1/collections/{collection_id}/items/{chapter['id']}/move",
headers=headers,
json={"parent_id": section["id"], "position": 0},
)
assert cycle.status_code == 400
removed = await client.delete(
f"/api/v1/collections/{collection_id}/groups/{chapter['id']}", headers=headers
)
assert removed.json()["promoted_items"] == 1
detail = (await client.get(f"/api/v1/collections/{collection_id}", headers=headers)).json()
assert detail["items"][0]["name"] == "第一节"
assert detail["videos"][0]["collection_path"] == ["第一节"]
asyncio.run(scenario())
def test_webdav_upload_move_browse_range_and_delete_guard(tmp_path: Path):
app, media, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic "
+ base64.b64encode(f"imagefind:{token}".encode()).decode()
}
payload = b"private-webdav-video"
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
configured = await client.patch(
"/api/v1/webdav/config",
headers={
**bearer,
"X-Forwarded-Proto": "http",
"X-Forwarded-Host": "192.168.5.100:80",
"X-Forwarded-Port": "5666",
},
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
assert configured.status_code == 200
assert configured.json()["url"] == "http://192.168.5.100:8765/webdav/"
assert configured.json()["direct_access"] is True
assert configured.json()["direct_port"] == 8765
assert configured.json()["gateway_supported"] is False
unauthenticated = await client.request("PROPFIND", "/webdav/", headers={"Depth": "0"})
assert unauthenticated.status_code == 401
assert unauthenticated.headers["www-authenticate"].startswith("Basic realm=")
assert (await client.request("MKCOL", "/webdav/旅行", headers=basic)).status_code == 201
uploaded = await client.put(
"/webdav/旅行/movie.mp4.part", headers=basic, content=payload
)
assert uploaded.status_code == 201
listing = await client.request(
"PROPFIND", "/webdav/旅行/", headers={**basic, "Depth": "1"}
)
assert listing.status_code == 207
assert "movie.mp4.part" in listing.text
moved = await client.request(
"MOVE",
"/webdav/旅行/movie.mp4.part",
headers={
**basic,
"Destination": "http://test/webdav/%E6%97%85%E8%A1%8C/movie.mp4",
},
)
assert moved.status_code == 201
upload = app.state.services.uploads.list()[0]
assert upload["origin"] == "webdav"
assert upload["collection_id"]
internal = app.state.services.uploads._get(upload["id"])
assert internal["webdav_path"] == "旅行/movie.mp4"
assert internal["content_sha256_verified"] is True
pending = await client.request(
"HEAD", "/webdav/旅行/movie.mp4", headers=basic
)
assert pending.status_code == 200
assert pending.headers["x-imagefind-upload-state"] == "accepted"
pending_listing = await client.request(
"PROPFIND", "/webdav/旅行/", headers={**basic, "Depth": "1"}
)
assert "movie.mp4" in pending_listing.text
app.state.services.uploads.transfer(internal["job_id"], upload["id"])
accepted = await client.request(
"HEAD", "/webdav/旅行/movie.mp4", headers=basic
)
assert accepted.status_code == 200
assert accepted.headers["retry-after"] == "2"
unavailable = await client.get("/webdav/旅行/movie.mp4", headers=basic)
assert unavailable.status_code == 503
assert unavailable.headers["retry-after"] == "2"
assert int(unavailable.headers["content-length"]) == len(unavailable.content)
assert int(unavailable.headers["content-length"]) != len(payload)
app.state.services.scanner.refresh_path(
"refresh", source_id, "dav/旅行/movie.mp4", upload["id"]
)
assert (media / "dav" / "旅行" / "movie.mp4").read_bytes() == payload
ranged = await client.get(
"/webdav/旅行/movie.mp4", headers={**basic, "Range": "bytes=0-6"}
)
assert ranged.status_code == 206
assert ranged.content == payload[:7]
head = await client.request(
"HEAD", "/webdav/旅行/movie.mp4", headers={**basic, "Range": "bytes=0-6"}
)
assert head.status_code == 206
assert head.headers["content-length"] == "7"
assert head.headers["content-range"] == f"bytes 0-6/{len(payload)}"
assert head.content == b""
guarded = await client.request("DELETE", "/webdav/旅行/movie.mp4", headers=basic)
assert guarded.status_code == 405
with app.state.services.db.read() as conn:
video_id = conn.execute(
"SELECT id FROM videos WHERE source_id=? AND source_key=?",
(source_id, "dav/旅行/movie.mp4"),
).fetchone()["id"]
deleted = await client.delete(
f"/api/v1/videos/{video_id}?delete_source=true", headers=bearer
)
assert deleted.status_code == 200
assert app.state.services.uploads._get(upload["id"])["status"] == "completed"
after_delete = await client.request(
"PROPFIND", "/webdav/旅行/", headers={**basic, "Depth": "1"}
)
assert after_delete.status_code == 207
assert "movie.mp4" not in after_delete.text
assert (await client.head("/webdav/旅行/movie.mp4", headers=basic)).status_code == 404
assert (await client.get("/webdav/旅行/movie.mp4", headers=basic)).status_code == 404
asyncio.run(scenario())
def test_webdav_maps_nested_directories_to_collection_groups(tmp_path: Path):
app, media, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
assert (await client.request("MKCOL", "/webdav/课程", headers=basic)).status_code == 201
assert (
await client.request("MKCOL", "/webdav/课程/第一章", headers=basic)
).status_code == 201
assert (
await client.request("MKCOL", "/webdav/课程/第一章/第一节", headers=basic)
).status_code == 201
uploaded = await client.put(
"/webdav/课程/第一章/第一节/clip.mp4", headers=basic, content=b"nested-video"
)
assert uploaded.status_code == 201
upload = app.state.services.uploads.list()[0]
assert upload["collection_parent_id"]
internal = app.state.services.uploads._get(upload["id"])
app.state.services.uploads.transfer(internal["job_id"], upload["id"])
app.state.services.scanner.refresh_path(
"refresh", source_id, "dav/课程/第一章/第一节/clip.mp4", upload["id"]
)
assert (media / "dav" / "课程" / "第一章" / "第一节" / "clip.mp4").is_file()
listing = await client.request(
"PROPFIND", "/webdav/课程/第一章/第一节/", headers={**basic, "Depth": "1"}
)
assert listing.status_code == 207
assert "clip.mp4" in listing.text
collections = (await client.get("/api/v1/collections", headers=bearer)).json()
detail = (
await client.get(f"/api/v1/collections/{collections[0]['id']}", headers=bearer)
).json()
assert detail["videos"][0]["collection_path"] == ["第一章", "第一节"]
asyncio.run(scenario())
def test_webdav_root_upload_is_not_added_to_a_collection(tmp_path: Path, monkeypatch):
app, media, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
response = await client.put(
"/webdav/root.mp4", headers=basic, content=b"root-video"
)
assert response.status_code == 201
upload = app.state.services.uploads.list()[0]
assert upload["collection_id"] is None
internal = app.state.services.uploads._get(upload["id"])
assert internal["webdav_path"] == "root.mp4"
assert internal["content_sha256"] == hashlib.sha256(b"root-video").hexdigest()
assert internal["content_sha256_verified"] is True
assert (await client.get("/api/v1/videos", headers=bearer)).json() == []
def unexpected_rescan(*_args, **_kwargs):
raise AssertionError("verified WebDAV PUT must not be hashed again during transfer")
monkeypatch.setattr(app.state.services.uploads, "_content_sha256", unexpected_rescan)
app.state.services.uploads.transfer(internal["job_id"], upload["id"])
assert (await client.get("/api/v1/videos", headers=bearer)).json() == []
app.state.services.scanner.refresh_path(
"refresh", source_id, "dav/root.mp4", upload["id"]
)
assert (media / "dav" / "root.mp4").read_bytes() == b"root-video"
videos = (await client.get("/api/v1/videos", headers=bearer)).json()
assert len(videos) == 1
assert videos[0]["collection_id"] is None
listing = await client.request(
"PROPFIND", "/webdav/", headers={**basic, "Depth": "1"}
)
assert listing.status_code == 207
assert "root.mp4" in listing.text
asyncio.run(scenario())
def test_webdav_content_range_resumes_and_reports_offset(tmp_path: Path):
app, media, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
first = await client.put(
"/webdav/课程/章节/clip.mp4",
headers={**basic, "Content-Range": "bytes 0-3/8"},
content=b"abcd",
)
assert first.status_code == 204
assert first.headers["upload-offset"] == "4"
assert app.state.services.uploads.list() == []
head = await client.request(
"HEAD", "/webdav/课程/章节/clip.mp4", headers=basic
)
assert head.status_code == 200
assert head.headers["upload-offset"] == "4"
assert head.headers["upload-length"] == "8"
assert head.headers["x-imagefind-upload-state"] == "partial"
wrong = await client.put(
"/webdav/课程/章节/clip.mp4",
headers={**basic, "Content-Range": "bytes 3-7/8"},
content=b"defgh",
)
assert wrong.status_code == 409
assert wrong.headers["upload-offset"] == "4"
second = await client.put(
"/webdav/课程/章节/clip.mp4",
headers={**basic, "Content-Range": "bytes 4-7/8"},
content=b"efgh",
)
assert second.status_code == 204
upload = app.state.services.uploads.list()[0]
assert upload["collection_id"]
internal = app.state.services.uploads._get(upload["id"])
assert internal["content_sha256"] is None
assert internal["content_sha256_verified"] is False
app.state.services.uploads.transfer(internal["job_id"], upload["id"])
assert app.state.services.uploads._get(upload["id"])["content_sha256_verified"] is True
app.state.services.scanner.refresh_path(
"refresh", source_id, "dav/课程/章节/clip.mp4", upload["id"]
)
assert (media / "dav" / "课程" / "章节" / "clip.mp4").read_bytes() == b"abcdefgh"
asyncio.run(scenario())
def test_webdav_full_retry_without_checksum_reuses_active_upload(tmp_path: Path):
app, _, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
payload = b"response-was-lost"
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
first = await client.put("/webdav/retry.mp4", headers=basic, content=payload)
assert first.status_code == 201
second = await client.put("/webdav/retry.mp4", headers=basic, content=payload)
assert second.status_code == 204
assert second.headers["x-imagefind-deduplicated"] == "true"
assert second.headers["x-imagefind-upload-id"] == first.headers["x-imagefind-upload-id"]
assert len(app.state.services.uploads.list()) == 1
with app.state.services.db.read() as conn:
assert conn.execute("SELECT count(*) FROM webdav_staging").fetchone()[0] == 0
asyncio.run(scenario())
def test_webdav_same_size_different_content_is_not_deduplicated(tmp_path: Path):
app, _, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
first = await client.put("/webdav/changed.mp4", headers=basic, content=b"first")
second = await client.put("/webdav/changed.mp4", headers=basic, content=b"other")
assert first.status_code == 201
assert second.status_code in {201, 204}
assert second.headers.get("x-imagefind-deduplicated") is None
assert len(app.state.services.uploads.list()) == 2
asyncio.run(scenario())
def test_webdav_head_repairs_delayed_offset_and_locked_put_reports_it(tmp_path: Path):
app, _, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
virtual = "course/offset.mp4"
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
partial = await client.put(
f"/webdav/{virtual}",
headers={**basic, "Content-Range": "bytes 0-3/8"},
content=b"abcd",
)
assert partial.status_code == 204
with app.state.services.db.transaction() as conn:
conn.execute(
"UPDATE webdav_staging SET received_bytes=1,size_bytes=1 WHERE virtual_path=?",
(virtual,),
)
head = await client.head(f"/webdav/{virtual}", headers=basic)
assert head.status_code == 200
assert head.headers["upload-offset"] == "4"
lock = _path_lock(virtual)
assert lock.acquire(blocking=False)
try:
locked = await client.put(
f"/webdav/{virtual}",
headers={**basic, "Content-Range": "bytes 4-7/8"},
content=b"efgh",
)
finally:
lock.release()
assert locked.status_code == 423
assert locked.headers["upload-offset"] == "4"
assert locked.headers["retry-after"] == "2"
asyncio.run(scenario())
def test_webdav_same_path_sha256_uses_safe_instant_deduplication(tmp_path: Path):
app, _, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
payload = b"deduplicated-video"
digest = hashlib.sha256(payload).hexdigest()
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
first = await client.put("/webdav/same.mp4", headers=basic, content=payload)
assert first.status_code == 201
upload = app.state.services.uploads.list()[0]
internal = app.state.services.uploads._get(upload["id"])
app.state.services.uploads.transfer(internal["job_id"], upload["id"])
app.state.services.scanner.refresh_path(
"refresh", source_id, "dav/same.mp4", upload["id"]
)
duplicate = await client.put(
"/webdav/same.mp4",
headers={**basic, "X-Content-SHA256": digest},
content=payload,
)
assert duplicate.status_code == 204
assert duplicate.headers["x-imagefind-deduplicated"] == "true"
assert len(app.state.services.uploads.list()) == 1
asyncio.run(scenario())
def test_webdav_buffer_flush_does_not_create_sparse_offset(tmp_path: Path, monkeypatch):
app, _, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
chunk_size = 8 * 1024**2
payload_size = chunk_size + 1024**2
class Clock:
values = iter((0.0, 3.0, 4.0))
@classmethod
def monotonic(cls):
return next(cls.values, 4.0)
monkeypatch.setattr(webdav_module, "time", Clock)
async def body():
yield b"a" * chunk_size
yield b"b" * (payload_size - chunk_size)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
response = await client.put(
"/webdav/buffer.mp4",
headers={**basic, "Content-Length": str(payload_size)},
content=body(),
)
assert response.status_code == 201
upload = app.state.services.uploads.list()[0]
assert upload["size_bytes"] == payload_size
internal = app.state.services.uploads._get(upload["id"])
assert Path(internal["temp_path"]).stat().st_size == payload_size
assert Path(internal["temp_path"]).read_bytes() == b"a" * chunk_size + b"b" * (
payload_size - chunk_size
)
asyncio.run(scenario())
def test_webdav_resume_truncates_unconfirmed_tail_and_rejects_short_file(tmp_path: Path):
staging = tmp_path / "resume.part"
staging.write_bytes(b"confirmed-unconfirmed")
descriptor = webdav_module._open_staging(staging, len(b"confirmed"))
try:
assert webdav_module.os.fstat(descriptor).st_size == len(b"confirmed")
finally:
webdav_module.os.close(descriptor)
assert staging.read_bytes() == b"confirmed"
with pytest.raises(OSError, match="短于已确认续传偏移"):
webdav_module._open_staging(staging, len(b"confirmed") + 1)
def test_webdav_stream_does_not_issue_explicit_burst_disk_flush(tmp_path: Path, monkeypatch):
app, _, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
sync_calls: list[int] = []
def unexpected_fsync(file_descriptor: int) -> None:
sync_calls.append(webdav_module.os.fstat(file_descriptor).st_size)
raise AssertionError("WebDAV must not synchronously fsync in the PUT request")
def unexpected_fdatasync(_file_descriptor: int) -> None:
raise AssertionError("WebDAV must not synchronously fdatasync while receiving the body")
monkeypatch.setattr(webdav_module.os, "fsync", unexpected_fsync)
monkeypatch.setattr(webdav_module.os, "fdatasync", unexpected_fdatasync)
open_flags: list[int] = []
real_open = webdav_module.os.open
def recording_open(path, flags, mode=0o777):
open_flags.append(flags)
return real_open(path, flags, mode)
write_sizes: list[int] = []
writeback_ranges: list[tuple[int, int]] = []
real_pwrite_all = webdav_module._pwrite_all
def recording_pwrite_all(file_descriptor: int, data: bytes, offset: int) -> int:
write_sizes.append(len(data))
return real_pwrite_all(file_descriptor, data, offset)
def recording_writeback(_file_descriptor: int, offset: int, length: int) -> bool:
writeback_ranges.append((offset, length))
return True
monkeypatch.setattr(webdav_module.os, "open", recording_open)
monkeypatch.setattr(webdav_module, "_pwrite_all", recording_pwrite_all)
monkeypatch.setattr(webdav_module, "_queue_writeback", recording_writeback)
payload_size = 9 * 1024**2
async def body():
yield b"a" * (8 * 1024**2)
assert sync_calls == []
yield b"b" * (payload_size - 8 * 1024**2)
assert sync_calls == []
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
response = await client.put(
"/webdav/sync-once.mp4",
headers={**basic, "Content-Length": str(payload_size)},
content=body(),
)
assert response.status_code == 201
asyncio.run(scenario())
assert sync_calls == []
assert open_flags
assert not any(flags & getattr(webdav_module.os, "O_DSYNC", 0) for flags in open_flags)
assert write_sizes
assert max(write_sizes) <= webdav_module.WEBDAV_WRITE_QUANTUM
assert writeback_ranges
assert sum(length for _, length in writeback_ranges) == payload_size
def test_webdav_propfind_does_not_block_other_requests(tmp_path: Path, monkeypatch):
if sys.version_info >= (3, 13):
pytest.skip("Python 3.13 sandbox thread selector deadlock; fnOS ships Python 3.12")
app, _, source_id, token = _app(tmp_path)
bearer = {"Authorization": f"Bearer {token}"}
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode()
}
async def threaded(_app, function, /, *args, **kwargs):
return await asyncio.to_thread(function, *args, **kwargs)
monkeypatch.setattr(webdav_module, "_background_io", threaded)
def slow_propfind(*_args, **_kwargs):
time.sleep(0.25)
return webdav_module.Response(status_code=207)
monkeypatch.setattr(webdav_module, "_propfind", slow_propfind)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": "dav"},
)
request = asyncio.create_task(
client.request("PROPFIND", "/webdav/", headers={**basic, "Depth": "0"})
)
await asyncio.sleep(0.02)
started = time.monotonic()
status = await client.get("/api/v1/status")
elapsed = time.monotonic() - started
assert status.status_code == 200
assert elapsed < 0.15
assert (await request).status_code == 207
asyncio.run(scenario())
def test_webdav_cannot_be_enabled_without_native_direct_access(tmp_path: Path):
app, _, source_id, token = _app(tmp_path, direct_access=False, port=9876)
bearer = {"Authorization": f"Bearer {token}"}
async def scenario():
transport = httpx.ASGITransport(app=app, root_path="/app/imagefind")
async with httpx.AsyncClient(
transport=transport, base_url="http://192.168.5.100:5666"
) as client:
config = await client.get("/api/v1/webdav/config", headers=bearer)
assert config.status_code == 200
assert config.json()["url"] == "http://192.168.5.100:9876/webdav/"
assert config.json()["direct_access"] is False
assert config.json()["gateway_supported"] is False
enabled = await client.patch(
"/api/v1/webdav/config",
headers=bearer,
json={"enabled": True, "source_id": source_id, "relative_path": ""},
)
assert enabled.status_code == 409
assert "直接 Web/API 访问" in enabled.json()["detail"]
asyncio.run(scenario())
def test_webdav_propfind_hrefs_include_gateway_root_path(tmp_path: Path):
app, _, source_id, token = _app(tmp_path)
app.state.services.db.set_setting(
"webdav_server", {"enabled": True, "source_id": source_id, "relative_path": ""}
)
basic = {
"Authorization": "Basic " + base64.b64encode(f"imagefind:{token}".encode()).decode(),
"Depth": "1",
}
async def scenario():
transport = httpx.ASGITransport(app=app, root_path="/app/imagefind")
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.request("PROPFIND", "/webdav/", headers=basic)
assert response.status_code == 207
assert "/app/imagefind/webdav/" in response.text
asyncio.run(scenario())
File diff suppressed because it is too large Load Diff
+226
View File
@@ -0,0 +1,226 @@
from __future__ import annotations
from pathlib import Path
from types import SimpleNamespace
import pytest
from imagefind.config import Settings
from imagefind.resources import ResourceGovernor
from imagefind.scanner import Scanner
from imagefind.speech_quality import transcript_quality_state
from imagefind.uploads import UploadService
from imagefind.usage import StorageUsageService
def _upload(**values) -> dict:
return {
"id": "upload",
"status": "completed",
"phase": "ai_queued",
"storage_backend": "openlist_native",
"size_bytes": 100,
"bytes_received": 100,
"progress": 1,
"external_progress": 1,
"external_status": "provider-internal-state",
"external_error": "provider-secret-error",
"external_task_id": "provider-task-12345678",
"deduplicated": 0,
**values,
}
def test_upload_public_contract_reports_ai_queue_without_leaking_provider_state():
result = UploadService._public(_upload())
assert result["stage"] == "ai_queued"
assert result["stage_label"] == "已加入媒体库,AI 已排队"
assert result["external_status"] == "completed"
assert "external_state" not in result
assert "external_error" not in result
assert "provider-internal-state" not in str(result)
assert "provider-secret-error" not in str(result)
assert result["external_task_ref"] == "12345678"
assert [stage["state"] for stage in result["stages"]] == ["completed"] * 4
assert result["stages"][-1]["label"] == "加入媒体库并排队 AI"
historical = UploadService._public(_upload(phase="cataloging"))
assert historical["stage"] == "ai_queued"
assert historical["stages"][-1]["state"] == "completed"
def test_native_restore_refresh_preserves_the_physical_object_mapping():
restored = Scanner._native_refresh_record(
None,
{
"physical_path": "cloud/library/opaque-name.bin",
"physical_size_bytes": 356_793,
"logical_size_bytes": 356_665,
},
)
assert restored == {
"storage_backend": "openlist_native",
"external_target_path": "cloud/library/opaque-name.bin",
"external_size_bytes": 356_793,
"size_bytes": 356_665,
"content_sha256": "",
}
assert Scanner._native_refresh_record(None, {"physical_path": "", "physical_size_bytes": 1}) is None
@pytest.mark.parametrize(
("score", "accepted", "rejected", "flags", "expected"),
[
(None, 0, 0, [], "empty"),
(0.82, 8, 2, [], "ready"),
(0.60, 5, 3, [], "filtered"),
(0.90, 8, 0, ["language_script_conflict"], "low_quality"),
(0.40, 2, 8, [], "low_quality"),
],
)
def test_transcript_quality_state_contract(score, accepted, rejected, flags, expected):
assert transcript_quality_state(score, accepted, rejected, flags) == expected
def test_running_lane_clears_stale_io_pressure_reason(monkeypatch, tmp_path: Path):
db = SimpleNamespace(setting=lambda _key, default=None: default)
governor = ResourceGovernor(db, Settings(data_dir=tmp_path))
governor._pressure_reasons["ai"] = "磁盘 I/O 持续繁忙,后台任务已让出资源"
monkeypatch.setattr(
governor,
"sample",
lambda: {
"cpu_percent": 10,
"memory_total_bytes": 8 * 1024**3,
"memory_available_bytes": 6 * 1024**3,
"disk_available_bytes": 100 * 1024**3,
"io": {"supported": False},
},
)
assert governor.pressure_reason(lane="ai", running=True) is None
assert governor._pressure_reasons["ai"] is None
class _Sources:
def __init__(self, source: dict, connector=None):
self.source = source
self._connector = connector
def get(self, source_id: str):
assert source_id == self.source["id"]
return self.source
def connector(self, source_id: str):
assert source_id == self.source["id"]
return self._connector
def _usage(source: dict, *, rclone=None, connector=None) -> StorageUsageService:
return StorageUsageService(
db=None,
settings=SimpleNamespace(),
rclone=rclone,
sources=_Sources(source, connector),
cache_seconds=30,
)
def test_local_capacity_uses_source_filesystem(monkeypatch, tmp_path: Path):
source = {"id": "local", "kind": "local", "config": {"path": str(tmp_path)}}
monkeypatch.setattr(
"imagefind.usage.shutil.disk_usage",
lambda path: SimpleNamespace(total=10_000, used=4_000, free=6_000)
if path == tmp_path
else None,
)
result = _usage(source).capacity("local")
assert (result["scope"], result["provider"], result["status"]) == (
"local",
"local",
"available",
)
assert result["available_bytes"] == 6_000
def test_openlist_capacity_uses_remote_quota_and_cache():
source = {
"id": "cloud",
"kind": "webdav",
"config": {"driver": "alist", "storage_backend": "openlist_native"},
}
calls = []
rclone = SimpleNamespace(
about=lambda value: calls.append(value["id"])
or {"total_bytes": 5_000, "used_bytes": 1_000, "available_bytes": 4_000}
)
usage = _usage(source, rclone=rclone)
first = usage.capacity("cloud")
second = usage.capacity("cloud")
assert first["scope"] == "remote"
assert first["provider"] == "rclone"
assert first["available_bytes"] == 4_000
assert second["available_bytes"] == 4_000
assert calls == ["cloud"]
def test_remote_capacity_manual_override_wins_without_querying_the_nas_or_provider():
total = 6 * 1024**4
available = 4 * 1024**4
source = {
"id": "cloud",
"kind": "webdav",
"config": {
"driver": "alist",
"storage_backend": "openlist_native",
"capacity_override_total_bytes": total,
"capacity_override_available_bytes": available,
},
}
rclone = SimpleNamespace(about=lambda _source: (_ for _ in ()).throw(AssertionError("must not query")))
result = _usage(source, rclone=rclone).capacity("cloud")
assert result["provider"] == "manual"
assert result["scope"] == "remote"
assert result["total_bytes"] == total
assert result["used_bytes"] == total - available
assert result["available_bytes"] == available
def test_webdav_capacity_supports_quota_and_distinguishes_unsupported_from_unreachable():
xml = b"""<?xml version='1.0'?>
<d:multistatus xmlns:d='DAV:'><d:response><d:propstat><d:prop>
<d:quota-used-bytes>300</d:quota-used-bytes>
<d:quota-available-bytes>700</d:quota-available-bytes>
</d:prop></d:propstat></d:response></d:multistatus>"""
class Connector:
base_url = "https://dav.example/root/"
def __init__(self, response):
self.client = SimpleNamespace(request=lambda *_args, **_kwargs: response)
self.closed = False
def close(self):
self.closed = True
source = {"id": "dav", "kind": "webdav", "config": {"driver": "webdav"}}
response = SimpleNamespace(status_code=207, content=xml)
connector = Connector(response)
available = _usage(source, connector=connector).capacity("dav")
assert available["provider"] == "webdav"
assert available["total_bytes"] == 1_000
assert connector.closed is True
no_quota = Connector(SimpleNamespace(status_code=207, content=b"<d:multistatus xmlns:d='DAV:'/>"))
unsupported = _usage(source, connector=no_quota).capacity("dav")
assert unsupported["status"] == "unsupported"
failed = Connector(SimpleNamespace(status_code=503, content=b""))
unreachable = _usage(source, connector=failed).capacity("dav")
assert unreachable["status"] == "unreachable"
assert unreachable["total_bytes"] == 0