Files
imagefind/scripts/live-upload-acceptance.py
T

236 lines
8.5 KiB
Python

#!/usr/bin/env python3
"""Run the resumable browser-upload acceptance path against a disposable live server.
The API token is read from a mode-0600 file and is never included in output. The
script only creates objects prefixed by the supplied run id and records their ids
so a later cleanup can be narrowly scoped.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import time
from pathlib import Path
from typing import Any
import httpx
def request(client: httpx.Client, method: str, path: str, **kwargs: Any) -> httpx.Response:
started = time.perf_counter()
response = client.request(method, path, **kwargs)
elapsed_ms = round((time.perf_counter() - started) * 1000, 1)
if response.is_error:
detail = response.text[:1000]
raise RuntimeError(f"{method} {path} -> {response.status_code} in {elapsed_ms} ms: {detail}")
return response
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
while block := handle.read(1024 * 1024):
digest.update(block)
return digest.hexdigest()
def ensure_metadata(client: httpx.Client, run_id: str) -> tuple[str, str, str]:
collection_name = f"{run_id}-合集"
collections = request(client, "GET", "/api/v1/collections").json()
collection = next((item for item in collections if item["name"] == collection_name), None)
if collection is None:
collection = request(
client,
"POST",
"/api/v1/collections",
json={"name": collection_name, "description": "ImageFind 自动验收临时合集"},
).json()
group_name = f"{run_id}-分类"
groups = request(client, "GET", "/api/v1/tag-groups").json()
group = next((item for item in groups if item["name"] == group_name), None)
if group is None:
group = request(
client,
"POST",
"/api/v1/tag-groups",
json={"name": group_name, "selection_mode": "multi", "sort_order": 9999},
).json()
group_id = group["id"]
else:
group_id = group["id"]
tag_name = f"{run_id}-样本"
tags = request(client, "GET", f"/api/v1/tags?group_id={group_id}").json()
tag = next((item for item in tags if item["name"] == tag_name), None)
if tag is None:
tag = request(
client,
"POST",
"/api/v1/tags",
json={"group_id": group_id, "name": tag_name},
).json()
return collection["id"], group_id, tag["id"]
def upload_one(
client: httpx.Client,
fixture: Path,
source_id: str,
collection_id: str,
tag_id: str,
run_id: str,
case_id: str,
) -> tuple[dict[str, Any], list[float]]:
case_prefix = f"-{case_id}" if case_id else ""
filename = f"{run_id}{case_prefix}-{fixture.name}"
existing = request(client, "GET", "/api/v1/uploads?limit=500").json()
upload = next(
(
item
for item in existing
if item.get("filename") == filename
and item.get("status") in {"receiving", "queued", "transferring", "indexing"}
),
None,
)
latencies: list[float] = []
if upload is None:
upload = request(
client,
"POST",
"/api/v1/uploads",
json={
"source_id": source_id,
"relative_path": "ingest",
"filename": filename,
"title": f"自动验收 · {fixture.stem}",
"collection_id": collection_id,
"tag_ids": [tag_id],
"size_bytes": fixture.stat().st_size,
"sha256": sha256(fixture),
"conflict": "skip",
},
).json()
upload_id = upload["id"]
if upload.get("status") == "completed":
return upload, latencies
hidden = request(client, "GET", "/api/v1/videos?limit=500").json()
if any(item.get("source_key", "").endswith(filename) for item in hidden):
raise AssertionError(f"unfinished upload leaked into /videos: {filename}")
chunk_size = int(upload["chunk_size"])
received = {int(index) for index in upload.get("received_chunks", upload.get("received", []))}
with fixture.open("rb") as handle:
for index in range(int(upload["total_chunks"])):
data = handle.read(chunk_size)
if index in received:
continue
started = time.perf_counter()
request(
client,
"PUT",
f"/api/v1/uploads/{upload_id}/chunks/{index}",
content=data,
headers={"X-Chunk-SHA256": hashlib.sha256(data).hexdigest()},
)
latencies.append(round((time.perf_counter() - started) * 1000, 1))
videos = request(client, "GET", "/api/v1/videos?limit=500").json()
if any(item.get("source_key", "").endswith(filename) for item in videos):
raise AssertionError(f"partially received upload leaked into /videos: {filename}")
upload = request(client, "POST", f"/api/v1/uploads/{upload_id}/complete").json()
deadline = time.monotonic() + 180
while upload.get("status") not in {"completed", "failed", "cancelled"}:
if time.monotonic() >= deadline:
raise TimeoutError(f"upload did not finish: {filename} status={upload.get('status')}")
time.sleep(1)
uploads = request(client, "GET", "/api/v1/uploads?limit=500").json()
upload = next(item for item in uploads if item["id"] == upload_id)
if upload["status"] != "completed":
raise AssertionError(f"upload failed: {filename}: {upload.get('message')}")
return upload, latencies
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--base-url", required=True)
auth = parser.add_mutually_exclusive_group(required=True)
auth.add_argument("--token-file", type=Path)
auth.add_argument("--password")
parser.add_argument("--run-dir", type=Path, required=True)
parser.add_argument("--case-id", default="")
parser.add_argument("--fixture", action="append", default=[])
args = parser.parse_args()
state_path = args.run_dir / "state.json"
state = json.loads(state_path.read_text())
headers = {}
if args.token_file is not None:
token = args.token_file.read_text().strip()
if not token:
raise RuntimeError("empty API token")
headers["Authorization"] = f"Bearer {token}"
client = httpx.Client(
base_url=args.base_url.rstrip("/"),
headers=headers,
timeout=httpx.Timeout(60, connect=10),
)
if args.password is not None:
login = request(
client,
"POST",
"/api/v1/auth/login",
json={"password": args.password, "remember_device": False},
).json()
csrf = str(login.get("csrf_token") or "")
if not csrf:
raise AssertionError("login response omitted CSRF token")
client.headers["X-CSRF-Token"] = csrf
collection_id, group_id, tag_id = ensure_metadata(client, state["run_id"])
state.update({"collection_id": collection_id, "tag_group_id": group_id, "tag_id": tag_id})
results = []
all_latencies: list[float] = []
fixture_names = args.fixture or ["e2e-positive.mp4", "e2e-negative.mp4", "e2e-fallback.mpg"]
for name in fixture_names:
upload, latencies = upload_one(
client,
args.run_dir / name,
state["source_id"],
collection_id,
tag_id,
state["run_id"],
args.case_id.strip(),
)
results.append(
{
"id": upload["id"],
"filename": upload["filename"],
"status": upload["status"],
"progress": upload["progress"],
"message": upload.get("message"),
}
)
all_latencies.extend(latencies)
state["upload_ids"] = [item["id"] for item in results]
state_path.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n")
report = {
"uploads": results,
"chunk_latency_ms": {
"count": len(all_latencies),
"max": max(all_latencies, default=0),
"average": round(sum(all_latencies) / len(all_latencies), 1) if all_latencies else 0,
},
}
(args.run_dir / "upload-report.json").write_text(
json.dumps(report, ensure_ascii=False, indent=2) + "\n"
)
print(json.dumps(report, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()