1072 lines
41 KiB
Python
1072 lines
41 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from imagefind import storage as storage_module
|
|
from imagefind import uploads as uploads_module
|
|
from imagefind.config import Settings
|
|
from imagefind.database import DatabaseTransientError, utcnow
|
|
from imagefind.jobs import JobRetry
|
|
from imagefind.main import create_app
|
|
from imagefind.remote import AlistClient, openlist_endpoint_from_webdav_url
|
|
from imagefind.storage import StorageService, TransferCancelled
|
|
from imagefind.uploads import UploadService
|
|
|
|
|
|
def _app(tmp_path: Path):
|
|
settings = Settings(
|
|
data_dir=tmp_path / "data",
|
|
embedding_backend="hash",
|
|
scan_interval_seconds=86400,
|
|
upload_chunk_mb=1,
|
|
upload_staging_gb=1,
|
|
upload_reserve_gb=0,
|
|
)
|
|
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, {"Authorization": f"Bearer {token}"}
|
|
|
|
|
|
def test_chunked_upload_transfer_refresh_and_download(tmp_path: Path):
|
|
app, media, source_id, headers = _app(tmp_path)
|
|
payload = b"private-video-data" * 4096
|
|
|
|
async def scenario():
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
created = await client.post(
|
|
"/api/v1/uploads",
|
|
headers=headers,
|
|
json={
|
|
"source_id": source_id,
|
|
"relative_path": "imports",
|
|
"filename": "ABC-123.mp4",
|
|
"title": " 我的自定义标题 ",
|
|
"size_bytes": len(payload),
|
|
},
|
|
)
|
|
assert created.status_code == 201
|
|
upload = created.json()
|
|
assert upload["title"] == "我的自定义标题"
|
|
assert upload["total_chunks"] == 1
|
|
digest = hashlib.sha256(payload).hexdigest()
|
|
chunk = await client.put(
|
|
f"/api/v1/uploads/{upload['id']}/chunks/0",
|
|
headers={**headers, "X-Chunk-SHA256": digest},
|
|
content=payload,
|
|
)
|
|
assert chunk.status_code == 200
|
|
completed = await client.post(f"/api/v1/uploads/{upload['id']}/complete", headers=headers)
|
|
assert completed.status_code == 200
|
|
job_id = completed.json()["job_id"]
|
|
|
|
app.state.services.uploads.transfer(job_id, upload["id"])
|
|
assert (media / "imports" / "ABC-123.mp4").read_bytes() == payload
|
|
app.state.services.scanner.refresh_path("refresh-test", source_id, "imports/ABC-123.mp4", upload["id"])
|
|
|
|
uploads = (await client.get("/api/v1/uploads", headers=headers)).json()
|
|
assert uploads[0]["status"] == "completed"
|
|
assert uploads[0]["title"] == "我的自定义标题"
|
|
videos = (await client.get("/api/v1/videos", headers=headers)).json()
|
|
assert videos[0]["catalog_code"] == "ABC-123"
|
|
assert videos[0]["title"] == "我的自定义标题"
|
|
download = await client.get(videos[0]["download_url"], headers=headers)
|
|
assert download.content == payload
|
|
assert "ABC-123.mp4" in download.headers["content-disposition"]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_transient_target_failure_requeues_without_losing_local_staging(tmp_path: Path, monkeypatch):
|
|
app, _, source_id, _ = _app(tmp_path)
|
|
payload = b"retryable-upload"
|
|
upload = app.state.services.uploads.create(source_id, "imports", "retry.mp4", len(payload))
|
|
app.state.services.uploads.receive_chunk(upload["id"], 0, payload)
|
|
completed = app.state.services.uploads.complete(upload["id"])
|
|
temporary = Path(app.state.services.uploads._get(upload["id"])["temp_path"])
|
|
monkeypatch.setattr(
|
|
app.state.services.storage,
|
|
"write_file",
|
|
lambda *_args, **_kwargs: (_ for _ in ()).throw(ConnectionError("connection reset")),
|
|
)
|
|
|
|
with pytest.raises(JobRetry) as retry:
|
|
app.state.services.uploads.transfer(completed["job_id"], upload["id"])
|
|
|
|
current = app.state.services.uploads._get(upload["id"])
|
|
assert retry.value.delay_seconds == 2
|
|
assert current["status"] == "queued"
|
|
assert current["retry_count"] == 1
|
|
assert current["next_retry_at"]
|
|
assert current["resume_mode"] == "restart"
|
|
assert temporary.read_bytes() == payload
|
|
|
|
|
|
def test_upload_title_can_be_updated_before_completion_and_is_validated(tmp_path: Path):
|
|
app, _, source_id, headers = _app(tmp_path)
|
|
|
|
async def scenario():
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
created = await client.post(
|
|
"/api/v1/uploads",
|
|
headers=headers,
|
|
json={"source_id": source_id, "filename": "movie.mp4", "title": "初始标题", "size_bytes": 4},
|
|
)
|
|
assert created.status_code == 201
|
|
upload_id = created.json()["id"]
|
|
updated = await client.patch(
|
|
f"/api/v1/uploads/{upload_id}", headers=headers, json={"title": " 更新标题 "}
|
|
)
|
|
assert updated.status_code == 200
|
|
assert updated.json()["title"] == "更新标题"
|
|
cleared = await client.patch(f"/api/v1/uploads/{upload_id}", headers=headers, json={"title": " "})
|
|
assert cleared.status_code == 200
|
|
assert cleared.json()["title"] is None
|
|
too_long = await client.patch(f"/api/v1/uploads/{upload_id}", headers=headers, json={"title": "x" * 501})
|
|
assert too_long.status_code == 422
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_upload_history_pagination_is_stable_and_keeps_global_status(tmp_path: Path):
|
|
app, _, source_id, headers = _app(tmp_path)
|
|
uploads = app.state.services.uploads
|
|
created = [uploads.create(source_id, "", f"movie-{index:02d}.mp4", 1) for index in range(23)]
|
|
upload_ids = [item["id"] for item in created]
|
|
same_created_at = "2026-08-04T00:00:00+00:00"
|
|
with app.state.services.db.transaction() as conn:
|
|
conn.execute("UPDATE uploads SET status='completed',created_at=?", (same_created_at,))
|
|
conn.execute("UPDATE uploads SET status='receiving' WHERE id=?", (upload_ids[0],))
|
|
conn.execute("UPDATE uploads SET status='queued' WHERE id=?", (upload_ids[1],))
|
|
conn.execute("UPDATE uploads SET status='failed',error='network error' WHERE id=?", (upload_ids[2],))
|
|
|
|
async def scenario():
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
first = (await client.get("/api/v1/uploads?page=1&page_size=10", headers=headers)).json()
|
|
second = (await client.get("/api/v1/uploads?page=2&page_size=10", headers=headers)).json()
|
|
clamped = (await client.get("/api/v1/uploads?page=99&page_size=10", headers=headers)).json()
|
|
legacy = (await client.get("/api/v1/uploads?limit=4", headers=headers)).json()
|
|
|
|
expected = sorted(upload_ids, reverse=True)
|
|
assert (first["page"], first["pages"], first["total"], len(first["items"])) == (1, 3, 23, 10)
|
|
assert [item["id"] for item in first["items"]] == expected[:10]
|
|
assert [item["id"] for item in second["items"]] == expected[10:20]
|
|
assert not ({item["id"] for item in first["items"]} & {item["id"] for item in second["items"]})
|
|
assert (clamped["page"], len(clamped["items"])) == (3, 3)
|
|
assert first["active_count"] == 2
|
|
assert first["failed_count"] == 1
|
|
assert {item["id"] for item in first["status_items"]} == set(upload_ids[:3])
|
|
assert isinstance(legacy, list) and len(legacy) == 4
|
|
assert uploads.cached_paginate(1, 10) == first
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_upload_missing_chunks_and_hash_rejection(tmp_path: Path):
|
|
app, _, source_id, headers = _app(tmp_path)
|
|
payload = b"x" * (1024 * 1024 + 3)
|
|
|
|
async def scenario():
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
upload = (
|
|
await client.post(
|
|
"/api/v1/uploads",
|
|
headers=headers,
|
|
json={"source_id": source_id, "filename": "movie.mp4", "size_bytes": len(payload)},
|
|
)
|
|
).json()
|
|
bad = await client.put(
|
|
f"/api/v1/uploads/{upload['id']}/chunks/0",
|
|
headers={**headers, "X-Chunk-SHA256": "0" * 64},
|
|
content=payload[: 1024 * 1024],
|
|
)
|
|
assert bad.status_code == 400
|
|
good = await client.put(
|
|
f"/api/v1/uploads/{upload['id']}/chunks/0",
|
|
headers=headers,
|
|
content=payload[: 1024 * 1024],
|
|
)
|
|
assert good.status_code == 200
|
|
incomplete = await client.post(f"/api/v1/uploads/{upload['id']}/complete", headers=headers)
|
|
assert incomplete.json() == {"completed": False, "missing_chunks": [1]}
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_upload_cancel_is_consistent_and_commit_window_returns_conflict(tmp_path: Path):
|
|
app, _, source_id, headers = _app(tmp_path)
|
|
upload = app.state.services.uploads.create(source_id, "", "movie.mp4", 4)
|
|
app.state.services.uploads.receive_chunk(upload["id"], 0, b"data")
|
|
job_id = app.state.services.uploads.complete(upload["id"])["job_id"]
|
|
app.state.services.uploads.cancel(upload["id"])
|
|
|
|
with app.state.services.db.read() as conn:
|
|
upload_state = conn.execute("SELECT status FROM uploads WHERE id=?", (upload["id"],)).fetchone()
|
|
job_state = conn.execute("SELECT status FROM jobs WHERE id=?", (job_id,)).fetchone()
|
|
assert upload_state["status"] == "cancelled"
|
|
assert job_state["status"] == "cancelled"
|
|
with app.state.services.db.read() as conn:
|
|
retained_path = Path(
|
|
conn.execute("SELECT temp_path FROM uploads WHERE id=?", (upload["id"],)).fetchone()[0]
|
|
)
|
|
assert retained_path.read_bytes() == b"data"
|
|
|
|
committed = app.state.services.uploads.create(source_id, "", "committed.mp4", 4)
|
|
with app.state.services.db.transaction() as conn:
|
|
conn.execute(
|
|
"UPDATE uploads SET status='transferring',failure_stage='commit',target_key='committed.mp4' WHERE id=?",
|
|
(committed["id"],),
|
|
)
|
|
|
|
async def scenario():
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
response = await client.delete(f"/api/v1/uploads/{committed['id']}", headers=headers)
|
|
assert response.status_code == 409
|
|
assert "已经写入" in response.json()["detail"]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_cancelled_upload_requires_explicit_discard_to_delete_recovery_copy(tmp_path: Path):
|
|
app, _, source_id, headers = _app(tmp_path)
|
|
uploads = app.state.services.uploads
|
|
upload = uploads.create(source_id, "", "recoverable.mp4", 4)
|
|
uploads.receive_chunk(upload["id"], 0, b"data")
|
|
with uploads.db.read() as conn:
|
|
staging = Path(conn.execute("SELECT temp_path FROM uploads WHERE id=?", (upload["id"],)).fetchone()[0])
|
|
uploads.cancel(upload["id"])
|
|
assert staging.is_file()
|
|
|
|
async def scenario():
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
response = await client.delete(
|
|
f"/api/v1/uploads/{upload['id']}/recovery", headers=headers
|
|
)
|
|
assert response.status_code == 204
|
|
|
|
asyncio.run(scenario())
|
|
assert not staging.exists()
|
|
with uploads.db.read() as conn:
|
|
assert conn.execute("SELECT 1 FROM uploads WHERE id=?", (upload["id"],)).fetchone() is None
|
|
|
|
|
|
def test_upload_catalog_failure_retries_without_retransmitting(tmp_path: Path, monkeypatch):
|
|
app, media, source_id, _ = _app(tmp_path)
|
|
uploads = app.state.services.uploads
|
|
upload = uploads.create(source_id, "imports", "movie.mp4", 4)
|
|
uploads.receive_chunk(upload["id"], 0, b"data")
|
|
transfer_job_id = uploads.complete(upload["id"])["job_id"]
|
|
enqueue = uploads.jobs.enqueue
|
|
|
|
def fail_refresh(kind, payload, **kwargs):
|
|
if kind == "refresh_path":
|
|
raise RuntimeError("catalog transaction unavailable")
|
|
return enqueue(kind, payload, **kwargs)
|
|
|
|
monkeypatch.setattr(uploads.jobs, "enqueue", fail_refresh)
|
|
with pytest.raises(RuntimeError, match="catalog transaction"):
|
|
uploads.transfer(transfer_job_id, upload["id"])
|
|
|
|
with uploads.db.read() as conn:
|
|
failed = dict(
|
|
conn.execute("SELECT status,failure_stage,target_key FROM uploads WHERE id=?", (upload["id"],)).fetchone()
|
|
)
|
|
assert failed == {
|
|
"status": "failed",
|
|
"failure_stage": "catalog",
|
|
"target_key": "imports/movie.mp4",
|
|
}
|
|
assert (media / "imports" / "movie.mp4").read_bytes() == b"data"
|
|
|
|
monkeypatch.setattr(uploads.jobs, "enqueue", enqueue)
|
|
refresh_job_id = uploads.retry(upload["id"])
|
|
with uploads.db.read() as conn:
|
|
job = dict(conn.execute("SELECT kind,payload_json FROM jobs WHERE id=?", (refresh_job_id,)).fetchone())
|
|
state = conn.execute("SELECT status FROM uploads WHERE id=?", (upload["id"],)).fetchone()[0]
|
|
assert job["kind"] == "refresh_path"
|
|
assert json.loads(job["payload_json"])["key"] == "imports/movie.mp4"
|
|
assert state == "indexing"
|
|
|
|
|
|
def test_upload_recovery_retires_stale_transfer_after_destination_commit(tmp_path: Path):
|
|
app, media, source_id, _ = _app(tmp_path)
|
|
uploads = app.state.services.uploads
|
|
upload = uploads.create(source_id, "imports", "movie.mp4", 4)
|
|
uploads.receive_chunk(upload["id"], 0, b"data")
|
|
transfer_job_id = uploads.complete(upload["id"])["job_id"]
|
|
with uploads.db.read() as conn:
|
|
staging = Path(conn.execute("SELECT temp_path FROM uploads WHERE id=?", (upload["id"],)).fetchone()[0])
|
|
target = media / "imports" / "movie.mp4"
|
|
target.parent.mkdir(parents=True)
|
|
staging.replace(target)
|
|
with uploads.db.transaction() as conn:
|
|
conn.execute(
|
|
"UPDATE uploads SET status='transferring',failure_stage='commit',target_key='imports/movie.mp4' WHERE id=?",
|
|
(upload["id"],),
|
|
)
|
|
conn.execute(
|
|
"UPDATE jobs SET status='running',lease_owner='old-process:transfer',heartbeat_at=? WHERE id=?",
|
|
(utcnow(), transfer_job_id),
|
|
)
|
|
|
|
assert uploads.jobs.recover_stale() == 1
|
|
assert uploads.recover_pending() == 1
|
|
|
|
with uploads.db.read() as conn:
|
|
stale = conn.execute("SELECT status FROM jobs WHERE id=?", (transfer_job_id,)).fetchone()[0]
|
|
current = dict(
|
|
conn.execute("SELECT status,job_id,target_key FROM uploads WHERE id=?", (upload["id"],)).fetchone()
|
|
)
|
|
refresh = conn.execute("SELECT kind,status FROM jobs WHERE id=?", (current["job_id"],)).fetchone()
|
|
assert stale == "cancelled"
|
|
assert current["status"] == "indexing"
|
|
assert current["target_key"] == "imports/movie.mp4"
|
|
assert dict(refresh) == {"kind": "refresh_path", "status": "queued"}
|
|
|
|
|
|
def test_upload_recovery_ignores_active_transfer_and_reconnects_failed_job(tmp_path: Path):
|
|
app, _, source_id, _ = _app(tmp_path)
|
|
uploads = app.state.services.uploads
|
|
upload = uploads.create(source_id, "imports", "movie.mp4", 4)
|
|
uploads.receive_chunk(upload["id"], 0, b"data")
|
|
transfer_job_id = uploads.complete(upload["id"])["job_id"]
|
|
|
|
assert uploads.recover_pending() == 0
|
|
with uploads.db.read() as conn:
|
|
active = dict(conn.execute("SELECT status,job_id FROM uploads WHERE id=?", (upload["id"],)).fetchone())
|
|
assert active == {"status": "queued", "job_id": transfer_job_id}
|
|
|
|
with uploads.db.transaction() as conn:
|
|
conn.execute(
|
|
"UPDATE jobs SET status='failed',error='deadlock detected',finished_at=? WHERE id=?",
|
|
(utcnow(), transfer_job_id),
|
|
)
|
|
conn.execute(
|
|
"UPDATE uploads SET status='transferring',progress=.86,message='传输到目标库 86%' WHERE id=?",
|
|
(upload["id"],),
|
|
)
|
|
|
|
assert uploads.recover_pending() == 1
|
|
with uploads.db.read() as conn:
|
|
recovered = dict(
|
|
conn.execute("SELECT status,job_id,message FROM uploads WHERE id=?", (upload["id"],)).fetchone()
|
|
)
|
|
replacement = dict(conn.execute("SELECT kind,status FROM jobs WHERE id=?", (recovered["job_id"],)).fetchone())
|
|
assert recovered["status"] == "queued"
|
|
assert recovered["job_id"] != transfer_job_id
|
|
assert "任务中断" in recovered["message"]
|
|
assert replacement == {"kind": "transfer_upload", "status": "queued"}
|
|
|
|
|
|
def test_locked_transfer_failure_does_not_leave_upload_transferring(tmp_path: Path, monkeypatch):
|
|
app, _, source_id, _ = _app(tmp_path)
|
|
uploads = app.state.services.uploads
|
|
upload = uploads.create(source_id, "imports", "movie.mp4", 4)
|
|
uploads.receive_chunk(upload["id"], 0, b"data")
|
|
transfer_job_id = uploads.complete(upload["id"])["job_id"]
|
|
original_get = uploads._get
|
|
reads = 0
|
|
|
|
def read_once(upload_id: str):
|
|
nonlocal reads
|
|
reads += 1
|
|
if reads > 1:
|
|
raise DatabaseTransientError("deadlock detected")
|
|
return original_get(upload_id)
|
|
|
|
monkeypatch.setattr(uploads, "_get", read_once)
|
|
monkeypatch.setattr(
|
|
uploads.storage,
|
|
"write_file",
|
|
lambda *_args, **_kwargs: (_ for _ in ()).throw(DatabaseTransientError("deadlock detected")),
|
|
)
|
|
|
|
with pytest.raises(JobRetry, match="传输中断"):
|
|
uploads.transfer(transfer_job_id, upload["id"])
|
|
|
|
with uploads.db.read() as conn:
|
|
state = dict(
|
|
conn.execute("SELECT status,retry_count,failure_stage FROM uploads WHERE id=?", (upload["id"],)).fetchone()
|
|
)
|
|
assert reads == 1
|
|
assert state == {"status": "queued", "retry_count": 1, "failure_stage": "transfer"}
|
|
|
|
|
|
def test_webdav_upload_uses_temporary_put_atomic_move_and_safe_cleanup(tmp_path: Path, monkeypatch):
|
|
app, _, _, _ = _app(tmp_path)
|
|
storage = app.state.services.storage
|
|
source = {
|
|
"id": "remote",
|
|
"kind": "webdav",
|
|
"config": {"writable": True, "driver": "webdav"},
|
|
}
|
|
monkeypatch.setattr(storage.sources, "get", lambda _source_id: source)
|
|
|
|
class Client:
|
|
def __init__(self, move_status=201):
|
|
self.move_status = move_status
|
|
self.requests: list[tuple[str, str, dict]] = []
|
|
self.puts: list[tuple[str, bytes]] = []
|
|
self.deletes: list[str] = []
|
|
|
|
def request(self, method, url, headers=None, content=None, timeout=None):
|
|
self.requests.append((method, url, headers or {}))
|
|
if method == "PROPFIND":
|
|
return SimpleNamespace(status_code=404)
|
|
if method == "MOVE":
|
|
return SimpleNamespace(status_code=self.move_status)
|
|
return SimpleNamespace(status_code=201)
|
|
|
|
def put(self, url, headers=None, content=None, timeout=None):
|
|
self.puts.append((url, b"".join(content)))
|
|
return SimpleNamespace(status_code=201)
|
|
|
|
def delete(self, url):
|
|
self.deletes.append(url)
|
|
return SimpleNamespace(status_code=204)
|
|
|
|
class Connector:
|
|
base_url = "https://dav.example/media/"
|
|
|
|
def __init__(self, client):
|
|
self.client = client
|
|
self.closed = False
|
|
|
|
def url_for(self, key):
|
|
return self.base_url + key
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
payload = tmp_path / "remote-upload.part"
|
|
payload.write_bytes(b"remote-data")
|
|
client = Client()
|
|
connector = Connector(client)
|
|
monkeypatch.setattr(storage.sources, "connector", lambda _source_id: connector)
|
|
before: list[str] = []
|
|
key = storage.write_file(
|
|
"remote",
|
|
"imports",
|
|
"movie.mp4",
|
|
payload,
|
|
conflict="replace",
|
|
operation_id="upload-1",
|
|
before_commit=before.append,
|
|
)
|
|
temporary_url = "https://dav.example/media/imports/.movie.mp4.imagefind-upload-1.part"
|
|
target_url = "https://dav.example/media/imports/movie.mp4"
|
|
assert key == "imports/movie.mp4"
|
|
assert client.puts == [(temporary_url, b"remote-data")]
|
|
move = next(item for item in client.requests if item[0] == "MOVE")
|
|
assert move[1] == temporary_url
|
|
assert move[2]["Destination"] == target_url
|
|
assert before == ["imports/movie.mp4"]
|
|
assert connector.closed is True
|
|
|
|
failing_client = Client(move_status=405)
|
|
failing_connector = Connector(failing_client)
|
|
monkeypatch.setattr(storage, "_REMOTE_MOVE_RETRY_DELAYS", (0, 0, 0))
|
|
monkeypatch.setattr(storage.sources, "connector", lambda _source_id: failing_connector)
|
|
with pytest.raises(RuntimeError, match="不支持安全原子写入"):
|
|
storage.write_file(
|
|
"remote",
|
|
"imports",
|
|
"other.mp4",
|
|
payload,
|
|
conflict="replace",
|
|
operation_id="upload-2",
|
|
)
|
|
failed_temp = "https://dav.example/media/imports/.other.mp4.imagefind-upload-2.part"
|
|
assert failing_client.puts == [(failed_temp, b"remote-data")]
|
|
assert failing_client.deletes[-1] == failed_temp
|
|
assert all(url != "https://dav.example/media/imports/other.mp4" for url, _ in failing_client.puts)
|
|
|
|
|
|
def test_remote_move_reconciles_false_failure_retries_and_true_failure(monkeypatch):
|
|
storage = object.__new__(StorageService)
|
|
monkeypatch.setattr(storage, "_REMOTE_MOVE_RETRY_DELAYS", (0, 0, 0))
|
|
|
|
class Client:
|
|
def __init__(self, outcomes):
|
|
self.outcomes = iter(outcomes)
|
|
self.files = {"source.mp4"}
|
|
self.moves = 0
|
|
|
|
def request(self, method, url, headers=None, content=None, timeout=None):
|
|
key = url.removeprefix("https://dav.example/media/")
|
|
if method == "PROPFIND":
|
|
return SimpleNamespace(status_code=207 if key in self.files else 404)
|
|
assert method == "MOVE"
|
|
assert timeout == storage._REMOTE_MOVE_TIMEOUT_SECONDS
|
|
self.moves += 1
|
|
outcome = next(self.outcomes)
|
|
if outcome in {"success", "moved-500"}:
|
|
self.files.discard(key)
|
|
destination = headers["Destination"].removeprefix("https://dav.example/media/")
|
|
self.files.add(destination)
|
|
return SimpleNamespace(status_code=201 if outcome == "success" else 500)
|
|
|
|
class Connector:
|
|
base_url = "https://dav.example/media/"
|
|
|
|
def __init__(self, client):
|
|
self.client = client
|
|
|
|
def url_for(self, key):
|
|
return self.base_url + key
|
|
|
|
false_failure = Client(["moved-500"])
|
|
storage._remote_move(
|
|
Connector(false_failure),
|
|
"source.mp4",
|
|
"trash/source.mp4",
|
|
overwrite=False,
|
|
failure_label="远程移动失败",
|
|
)
|
|
assert false_failure.moves == 1
|
|
assert false_failure.files == {"trash/source.mp4"}
|
|
|
|
retried = Client(["failed", "success"])
|
|
storage._remote_move(
|
|
Connector(retried),
|
|
"source.mp4",
|
|
"trash/source.mp4",
|
|
overwrite=False,
|
|
failure_label="远程移动失败",
|
|
)
|
|
assert retried.moves == 2
|
|
assert retried.files == {"trash/source.mp4"}
|
|
|
|
failed = Client(["failed"] * storage._REMOTE_MOVE_ATTEMPTS)
|
|
with pytest.raises(RuntimeError, match="源文件仍存在，目标文件不存在"):
|
|
storage._remote_move(
|
|
Connector(failed),
|
|
"source.mp4",
|
|
"trash/source.mp4",
|
|
overwrite=False,
|
|
failure_label="远程移动失败",
|
|
)
|
|
assert failed.moves == storage._REMOTE_MOVE_ATTEMPTS
|
|
assert failed.files == {"source.mp4"}
|
|
|
|
|
|
def test_remote_move_timeout_reconciles_completed_destination(monkeypatch):
|
|
storage = object.__new__(StorageService)
|
|
monkeypatch.setattr(storage, "_REMOTE_MOVE_RETRY_DELAYS", (0, 0, 0))
|
|
|
|
class Client:
|
|
def __init__(self):
|
|
self.files = {"source.mp4"}
|
|
|
|
def request(self, method, url, headers=None, content=None, timeout=None):
|
|
key = url.removeprefix("https://dav.example/media/")
|
|
if method == "PROPFIND":
|
|
return SimpleNamespace(status_code=207 if key in self.files else 404)
|
|
assert method == "MOVE"
|
|
assert timeout == storage._REMOTE_MOVE_TIMEOUT_SECONDS
|
|
self.files.remove(key)
|
|
self.files.add(headers["Destination"].removeprefix("https://dav.example/media/"))
|
|
raise TimeoutError("response timeout")
|
|
|
|
class Connector:
|
|
base_url = "https://dav.example/media/"
|
|
|
|
def __init__(self):
|
|
self.client = Client()
|
|
|
|
def url_for(self, key):
|
|
return self.base_url + key
|
|
|
|
connector = Connector()
|
|
storage._remote_move(
|
|
connector,
|
|
"source.mp4",
|
|
"trash/source.mp4",
|
|
overwrite=False,
|
|
failure_label="远程移动失败",
|
|
)
|
|
assert connector.client.files == {"trash/source.mp4"}
|
|
|
|
|
|
def test_remote_request_wait_can_be_cancelled_promptly(monkeypatch):
|
|
storage = object.__new__(StorageService)
|
|
monkeypatch.setattr(storage, "_REMOTE_CANCEL_POLL_SECONDS", 0.001)
|
|
started = threading.Event()
|
|
closed = threading.Event()
|
|
|
|
def request():
|
|
started.set()
|
|
closed.wait(2)
|
|
return SimpleNamespace(status_code=201)
|
|
|
|
with pytest.raises(TransferCancelled, match="已取消"):
|
|
storage._run_remote_request(request, closed.set, started.is_set)
|
|
assert closed.is_set()
|
|
|
|
|
|
def test_alist_stream_upload_forwards_existing_sha256_without_rehashing():
|
|
captured: dict[str, object] = {}
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
captured["method"] = request.method
|
|
captured["path"] = request.url.path
|
|
captured["headers"] = dict(request.headers)
|
|
captured["body"] = request.read()
|
|
return httpx.Response(200, json={"code": 200, "message": "success", "data": {}})
|
|
|
|
client = AlistClient(
|
|
"https://openlist.example",
|
|
"admin",
|
|
"secret",
|
|
root_path="云盘/测试",
|
|
)
|
|
client.client.close()
|
|
client.client = httpx.Client(transport=httpx.MockTransport(handler))
|
|
client._token = "private-token"
|
|
try:
|
|
client.put_stream(
|
|
"临时/movie.part",
|
|
iter((b"video-", b"payload")),
|
|
size_bytes=13,
|
|
content_sha256="a" * 64,
|
|
timeout=30,
|
|
)
|
|
finally:
|
|
client.close()
|
|
|
|
headers = captured["headers"]
|
|
assert captured["method"] == "PUT"
|
|
assert captured["path"] == "/api/fs/put"
|
|
assert captured["body"] == b"video-payload"
|
|
assert headers["authorization"] == "private-token"
|
|
assert headers["file-path"] == "/%E4%BA%91%E7%9B%98/%E6%B5%8B%E8%AF%95/%E4%B8%B4%E6%97%B6/movie.part"
|
|
assert headers["x-file-sha256"] == "a" * 64
|
|
assert headers["as-task"] == "false"
|
|
|
|
|
|
def test_direct_alist_storage_uses_hash_aware_api_then_atomic_move(tmp_path: Path, monkeypatch):
|
|
source = {
|
|
"id": "alist",
|
|
"kind": "webdav",
|
|
"config": {
|
|
"driver": "alist",
|
|
"mode": "direct",
|
|
"base_url": "https://openlist.example",
|
|
"root_path": "media",
|
|
"username": "admin",
|
|
"verify_tls": True,
|
|
"writable": True,
|
|
},
|
|
"secrets": {"password": "secret"},
|
|
}
|
|
|
|
class WebClient:
|
|
def __init__(self):
|
|
self.moves = []
|
|
|
|
def request(self, method, url, headers=None, content=None, timeout=None):
|
|
if method == "PROPFIND":
|
|
return SimpleNamespace(status_code=404)
|
|
if method == "MKCOL":
|
|
return SimpleNamespace(status_code=201)
|
|
assert method == "MOVE"
|
|
self.moves.append((url, headers, timeout))
|
|
return SimpleNamespace(status_code=201)
|
|
|
|
def delete(self, _url):
|
|
return SimpleNamespace(status_code=204)
|
|
|
|
class Connector:
|
|
base_url = "https://openlist.example/dav/media/"
|
|
|
|
def __init__(self):
|
|
self.client = WebClient()
|
|
self.closed = False
|
|
|
|
def url_for(self, key):
|
|
return self.base_url + key
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
connector = Connector()
|
|
|
|
class Sources:
|
|
def get(self, _source_id):
|
|
return source
|
|
|
|
def connector(self, _source_id):
|
|
return connector
|
|
|
|
controls = []
|
|
|
|
class Control:
|
|
def __init__(self, *_args, **kwargs):
|
|
self.kwargs = kwargs
|
|
self.uploads = []
|
|
self.closed = False
|
|
controls.append(self)
|
|
|
|
def put_stream(self, path, content, **kwargs):
|
|
self.uploads.append((path, b"".join(content), kwargs))
|
|
return {}
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
monkeypatch.setattr(storage_module, "AlistClient", Control)
|
|
storage = StorageService(None, SimpleNamespace(remote_timeout_seconds=30), Sources())
|
|
payload = tmp_path / "payload.part"
|
|
payload.write_bytes(b"remote-data")
|
|
key = storage.write_file(
|
|
"alist",
|
|
"imports",
|
|
"movie.mp4",
|
|
payload,
|
|
conflict="replace",
|
|
operation_id="upload-id",
|
|
content_sha256="b" * 64,
|
|
)
|
|
|
|
assert key == "imports/movie.mp4"
|
|
assert len(controls) == 1
|
|
temporary = "imports/.movie.mp4.imagefind-upload-id.part"
|
|
assert controls[0].uploads[0][0] == temporary
|
|
assert controls[0].uploads[0][1] == b"remote-data"
|
|
assert controls[0].uploads[0][2]["content_sha256"] == "b" * 64
|
|
assert connector.client.moves[0][0].endswith(temporary)
|
|
assert connector.client.moves[0][1]["Destination"].endswith("imports/movie.mp4")
|
|
assert controls[0].closed is True
|
|
assert connector.closed is True
|
|
|
|
|
|
def test_generic_openlist_webdav_source_is_probed_before_hash_aware_upload(monkeypatch):
|
|
source = {
|
|
"id": "generic-openlist",
|
|
"kind": "webdav",
|
|
"config": {
|
|
"base_url": "http://openlist.example:5244/dav/yidongpan/test%20space/",
|
|
"username": "admin",
|
|
"verify_tls": False,
|
|
"writable": True,
|
|
},
|
|
"secrets": {"password": "secret"},
|
|
}
|
|
controls = []
|
|
|
|
class Control:
|
|
def __init__(self, base_url, username, password, **kwargs):
|
|
self.base_url = base_url
|
|
self.username = username
|
|
self.password = password
|
|
self.kwargs = kwargs
|
|
self.probes = 0
|
|
self.closed = False
|
|
controls.append(self)
|
|
|
|
def probe(self):
|
|
self.probes += 1
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
monkeypatch.setattr(storage_module, "AlistClient", Control)
|
|
storage = StorageService(
|
|
None,
|
|
SimpleNamespace(remote_timeout_seconds=30),
|
|
SimpleNamespace(),
|
|
)
|
|
first = storage._openlist_upload_client(source["id"], source)
|
|
second = storage._openlist_upload_client(source["id"], source)
|
|
|
|
assert first is controls[0]
|
|
assert second is controls[1]
|
|
assert controls[0].base_url == "http://openlist.example:5244"
|
|
assert controls[0].kwargs["root_path"] == "yidongpan/test space"
|
|
assert controls[0].probes == 1
|
|
assert controls[1].probes == 0
|
|
|
|
|
|
def test_non_openlist_or_failed_probe_keeps_standard_webdav(monkeypatch):
|
|
storage = StorageService(
|
|
None,
|
|
SimpleNamespace(remote_timeout_seconds=30),
|
|
SimpleNamespace(),
|
|
)
|
|
assert openlist_endpoint_from_webdav_url("https://dav.example/remote/files") is None
|
|
|
|
source = {
|
|
"config": {
|
|
"base_url": "https://dav.example/dav/media/",
|
|
"username": "user",
|
|
"verify_tls": True,
|
|
},
|
|
"secrets": {"password": "secret"},
|
|
}
|
|
|
|
class Control:
|
|
def __init__(self, *_args, **_kwargs):
|
|
self.closed = False
|
|
|
|
def probe(self):
|
|
raise RuntimeError("not OpenList")
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
monkeypatch.setattr(storage_module, "AlistClient", Control)
|
|
assert storage._openlist_upload_client("ordinary", source) is None
|
|
|
|
|
|
def test_upload_transfer_progress_uses_bounded_database_checkpoints(monkeypatch):
|
|
class Connection:
|
|
def __init__(self):
|
|
self.executions = []
|
|
|
|
def execute(self, sql, args):
|
|
self.executions.append((sql, args))
|
|
|
|
class Database:
|
|
def __init__(self):
|
|
self.transactions = []
|
|
|
|
@contextmanager
|
|
def transaction(self):
|
|
connection = Connection()
|
|
self.transactions.append(connection)
|
|
yield connection
|
|
|
|
def write_with_retry(self, operation, *, timeout_seconds):
|
|
assert timeout_seconds == 2
|
|
with self.transaction() as connection:
|
|
return operation(connection)
|
|
|
|
class Jobs:
|
|
def __init__(self):
|
|
self.checkpoints = []
|
|
|
|
def checkpoint(self, job_id, *, persist):
|
|
self.checkpoints.append((job_id, persist))
|
|
|
|
service = object.__new__(UploadService)
|
|
service.db = Database()
|
|
service.jobs = Jobs()
|
|
service._progress_guard = threading.Lock()
|
|
service._progress_state = {}
|
|
clock = iter((100.0, 105.0, 106.0, 107.0))
|
|
monkeypatch.setattr(uploads_module.time, "monotonic", lambda: next(clock))
|
|
total = 512 * 1024**2
|
|
|
|
service._progress("upload", "job", 1, total)
|
|
service._progress("upload", "job", 100 * 1024**2, total)
|
|
service._progress("upload", "job", 300 * 1024**2, total)
|
|
service._progress("upload", "job", total, total)
|
|
|
|
assert len(service.db.transactions) == 3
|
|
assert service.jobs.checkpoints == [("job", False)] * 3
|
|
final_upload_args = service.db.transactions[-1].executions[0][1]
|
|
assert final_upload_args[0] == 0.99
|
|
assert final_upload_args[2] == "数据已发送，等待目标库确认"
|
|
|
|
|
|
def test_trash_layout_preserves_basename_and_remote_restore_avoids_rename():
|
|
storage = object.__new__(StorageService)
|
|
trash_id = "5dc19861-4aa4-4dc7-9149-73395b52ebea"
|
|
trash_key = storage._trash_key(trash_id, "课程/第一节.mp4")
|
|
assert trash_key == f".imagefind-trash/{trash_id}/第一节.mp4"
|
|
assert storage._trash_container_key(trash_key) == f".imagefind-trash/{trash_id}"
|
|
assert storage._trash_container_key(f".imagefind-trash/{trash_id}-第一节.mp4") is None
|
|
|
|
class Client:
|
|
def __init__(self, files):
|
|
self.files = set(files)
|
|
|
|
def request(self, method, url, headers=None, content=None):
|
|
assert method == "PROPFIND"
|
|
key = url.removeprefix("https://dav.example/media/")
|
|
return SimpleNamespace(status_code=207 if key in self.files else 404)
|
|
|
|
class Connector:
|
|
base_url = "https://dav.example/media/"
|
|
|
|
def __init__(self, files):
|
|
self.client = Client(files)
|
|
|
|
def url_for(self, key):
|
|
return self.base_url + key
|
|
|
|
item = {"id": trash_id, "original_key": "课程/第一节.mp4"}
|
|
assert storage._remote_restore_target(Connector(set()), item) == "课程/第一节.mp4"
|
|
assert storage._remote_restore_target(
|
|
Connector({"课程/第一节.mp4", f"课程/已恢复-{trash_id}/第一节.mp4"}),
|
|
item,
|
|
) == f"课程/已恢复-{trash_id}-2/第一节.mp4"
|
|
|
|
|
|
def test_local_trash_restore_and_path_protection(tmp_path: Path):
|
|
app, media, source_id, _ = _app(tmp_path)
|
|
target = media / "folder" / "movie.mp4"
|
|
target.parent.mkdir()
|
|
target.write_bytes(b"movie")
|
|
trash_id = app.state.services.storage.trash(source_id, "folder/movie.mp4")
|
|
assert not target.exists()
|
|
restored = app.state.services.storage.restore(trash_id)
|
|
assert restored == "folder/movie.mp4"
|
|
assert target.read_bytes() == b"movie"
|
|
try:
|
|
app.state.services.storage.list_dir(source_id, "../")
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
raise AssertionError("directory traversal must be rejected")
|
|
|
|
|
|
def test_alist_direct_stream_returns_redirect(tmp_path: Path):
|
|
app, _, _, headers = _app(tmp_path)
|
|
service = app.state.services
|
|
now = utcnow()
|
|
with service.db.transaction() as conn:
|
|
conn.execute(
|
|
"INSERT INTO sources(id,kind,name,config_json,secret_blob,created_at,updated_at) VALUES(?,?,?,?,?,?,?)",
|
|
(
|
|
"alist-source",
|
|
"webdav",
|
|
"AList 直连",
|
|
json.dumps(
|
|
{
|
|
"driver": "alist",
|
|
"mode": "direct",
|
|
"base_url": "https://alist.example",
|
|
"root_path": "private",
|
|
"username": "user",
|
|
}
|
|
),
|
|
service.secrets.encrypt_json({"password": "password"}),
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO videos(id,source_id,source_key,display_name,location,fingerprint,status,available,"
|
|
"created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?)",
|
|
(
|
|
"alist-video",
|
|
"alist-source",
|
|
"movie.mp4",
|
|
"movie.mp4",
|
|
"https://alist.example/dav/private/movie.mp4",
|
|
"fingerprint",
|
|
"indexed",
|
|
1,
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
service.storage.direct_url = lambda source_id, key: "https://provider.example/signed/movie.mp4"
|
|
|
|
async def scenario():
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test", follow_redirects=False) as client:
|
|
response = await client.get("/api/v1/videos/alist-video/stream", headers=headers)
|
|
assert response.status_code == 302
|
|
assert response.headers["location"] == "https://provider.example/signed/movie.mp4"
|
|
assert response.headers["cache-control"] == "no-store"
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_alist_recovery_import_api(tmp_path: Path):
|
|
app, _, _, headers = _app(tmp_path)
|
|
captured = {}
|
|
|
|
def restore_alist(name, base_url, root_path, username, password, crypt_password, crypt_salt, *, verify_tls):
|
|
captured.update(
|
|
name=name,
|
|
base_url=base_url,
|
|
root_path=root_path,
|
|
username=username,
|
|
password=password,
|
|
crypt_password=crypt_password,
|
|
crypt_salt=crypt_salt,
|
|
verify_tls=verify_tls,
|
|
)
|
|
return "restored-source"
|
|
|
|
app.state.services.sources.restore_alist = restore_alist
|
|
|
|
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/sources/alist/restore",
|
|
headers=headers,
|
|
json={
|
|
"name": "恢复的保险库",
|
|
"password": "new-alist-password",
|
|
"verify_tls": False,
|
|
"recovery": {
|
|
"format": "imagefind-rclone-crypt-v1",
|
|
"base_url": "https://alist.example",
|
|
"root_path": "private/videos",
|
|
"username": "admin",
|
|
"crypt_password": "crypt-secret",
|
|
"crypt_salt": "crypt-salt",
|
|
},
|
|
},
|
|
)
|
|
assert response.status_code == 201
|
|
assert response.json()["id"] == "restored-source"
|
|
assert captured == {
|
|
"name": "恢复的保险库",
|
|
"base_url": "https://alist.example/",
|
|
"root_path": "private/videos",
|
|
"username": "admin",
|
|
"password": "new-alist-password",
|
|
"crypt_password": "crypt-secret",
|
|
"crypt_salt": "crypt-salt",
|
|
"verify_tls": False,
|
|
}
|
|
|
|
invalid = await client.post(
|
|
"/api/v1/sources/alist/restore",
|
|
headers=headers,
|
|
json={"password": "x", "recovery": {"format": "unknown"}},
|
|
)
|
|
assert invalid.status_code == 422
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_single_worker_claims_heavy_jobs_in_queue_order(tmp_path: Path):
|
|
app, _, source_id, _ = _app(tmp_path)
|
|
jobs = app.state.services.jobs
|
|
transfer_id = jobs.enqueue("transfer_upload", {"upload_id": "one"})
|
|
scan_id = jobs.enqueue("scan_source", {"source_id": source_id})
|
|
|
|
# Claim selection is tested independently of host resource pressure; the
|
|
# production worker still consults the governor before every claim.
|
|
jobs.governor = None
|
|
transfer = jobs._claim()
|
|
general = jobs._claim()
|
|
|
|
assert transfer and transfer[0] == transfer_id
|
|
assert general and general[0] == scan_id
|