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