feat: add ImageFind application and release pipelines
This commit is contained in:
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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
@@ -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())
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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'"
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
@@ -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")
|
||||
@@ -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 == {}
|
||||
@@ -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
@@ -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")
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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())
|
||||
@@ -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
@@ -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
|
||||
Reference in New Issue
Block a user