"""批量真实计算、并行上限、失败隔离、恢复和保留期限。""" from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone import io import json from pathlib import Path import threading import time import zipfile import cv2 import numpy as np import pytest from fastapi import FastAPI from fastapi.testclient import TestClient import config from api.algorithm_routes import router from algorithms.characteristic_curve_recognition import service from algorithms.characteristic_curve_recognition.batch import BatchManager BASE = "/api/algorithms/characteristic-curve-recognition" @pytest.fixture def client(tmp_path, monkeypatch): monkeypatch.setattr(config, "ALGORITHM_JOB_DIR", str(tmp_path / "jobs")) monkeypatch.setattr(config, "ALGORITHM_MAX_PARALLEL", 2) monkeypatch.setattr(config, "ALGORITHM_MAX_PENDING_IMAGES", 5) monkeypatch.setattr(service, "_slots", threading.BoundedSemaphore(2)) app = FastAPI() app.include_router(router, prefix="/api") with TestClient(app) as client: yield client @pytest.fixture def chart(): image = np.full((320, 420, 3), 255, np.uint8) cv2.line(image, (45, 25), (45, 275), (0, 0, 0), 2) cv2.line(image, (45, 275), (390, 275), (0, 0, 0), 2) for offset in (0, 30): points = np.array([[x, int(45 + offset + .0012 * (x - 50) ** 2)] for x in range(60, 380)], np.int32) cv2.polylines(image, [points], False, (0, 0, 0), 2) return cv2.imencode(".png", image)[1].tobytes() def wait_until(fn, timeout=15): deadline = time.monotonic() + timeout while time.monotonic() < deadline: result = fn() if result: return result time.sleep(.02) raise AssertionError("等待任务状态超时") def finished(client, url): def check(): response = client.get(url) assert response.status_code == 200 job = response.json() return job if job["status"] in {"completed", "partial_failed", "failed", "interrupted"} else None return wait_until(check, timeout=45) def test_real_parallel_batch_partial_failure_and_download(client, chart, monkeypatch): run = service.subprocess.run barrier = threading.Barrier(2) lock = threading.Lock() active = peak = 0 def measured_run(*args, **kwargs): nonlocal active, peak with lock: active += 1 peak = max(peak, active) try: barrier.wait(timeout=10) return run(*args, **kwargs) finally: with lock: active -= 1 monkeypatch.setattr(service.subprocess, "run", measured_run) image = cv2.imdecode(np.frombuffer(chart, np.uint8), cv2.IMREAD_COLOR) jpeg = cv2.imencode(".jpg", image)[1].tobytes() response = client.post(BASE + "/batches", files=[ ("files", ("../../同名.png", chart)), ("files", ("同名.png", jpeg)), ("files", ("坏图.png", b"bad"))]) assert response.status_code == 202, response.text job = finished(client, response.headers["location"]) assert peak == 2 # 真正进入了两个并行子进程调用。 assert job["status"] == "partial_failed" assert job["progress"] == {"total": 3, "queued": 0, "running": 0, "completed": 2, "failed": 1, "interrupted": 0, "finished": 3} assert [item["filename"] for item in job["items"]] == ["同名.png", "同名.png", "坏图.png"] assert job["items"][2]["error"]["code"] == "invalid_image" for item in job["items"][:2]: result = client.get(item["result_url"]) assert result.status_code == 200 assert len(result.json()["model"]["curves"]) == 2 assert result.json()["geometry"]["grid"]["cells"] response = client.get(job["download_url"]) assert response.status_code == 200 with zipfile.ZipFile(io.BytesIO(response.content)) as archive: assert json.loads(archive.read("summary.json"))["status"] == "partial_failed" assert "items/0001/index.html" in archive.namelist() assert "items/0002/result.json" in archive.namelist() assert not any(name.startswith("items/0003") for name in archive.namelist()) root = Path(config.ALGORITHM_JOB_DIR) / job["id"] assert not list((root / "inputs").iterdir()) def test_queue_limits_shared_slots_and_progress(client, chart, monkeypatch): release = threading.Event() started = threading.Barrier(3) run_count = 0 lock = threading.Lock() def blocked_run(command, **kwargs): nonlocal run_count with lock: run_count += 1 first_pair = run_count <= 2 stage_path = Path(command[-1]) stage_path.write_text('{"stage":"piecewise_c1_fit"}') if first_pair: started.wait(timeout=5) assert release.wait(10) # 通过真实 service 的闸与超时转换，模拟快速成功产物。 root = Path(command[5]) (root / "diagnostics").mkdir(parents=True) for path in (root / "final_model.json", root / "diagnostics/geometry.json", root / "diagnostics/timing.json"): path.write_text("{}") return type("Result", (), {"returncode": 0})() monkeypatch.setattr(service.subprocess, "run", blocked_run) try: first = client.post(BASE + "/batches", files=[("files", ("x.png", chart))] * 3) assert first.status_code == 202 started.wait(timeout=5) second = client.post(BASE + "/batches", files=[("files", ("y.png", chart))] * 2) assert second.status_code == 202 excess = client.post(BASE + "/batches", files=[("files", ("z.png", chart))]) assert excess.status_code == 429 assert excess.json()["detail"]["code"] == "algorithm_queue_full" assert excess.headers["retry-after"] == "5" single = client.post(BASE, files={"file": ("single.png", chart)}) assert single.status_code == 429 job = client.get(first.headers["location"]).json() assert job["progress"]["running"] == 2 assert job["progress"]["queued"] == 1 assert job["items"][0]["stage"] == "piecewise_c1_fit" assert client.get(job["download_url"]).status_code == 409 assert client.get(job["status_url"] + "/items/0001/result").status_code == 409 assert client.get("/health-does-not-exist").status_code == 404 # HTTP 仍可响应。 finally: release.set() assert finished(client, first.headers["location"])["status"] == "completed" assert finished(client, second.headers["location"])["status"] == "completed" def test_timeout_isolated_and_queue_reusable(client, chart, monkeypatch): def timeout(*args, **kwargs): raise service.subprocess.TimeoutExpired(args[0], kwargs["timeout"]) monkeypatch.setattr(service.subprocess, "run", timeout) for _ in range(2): response = client.post(BASE + "/batches", files=[("files", ("x.png", chart))] * 2) assert response.status_code == 202 job = finished(client, response.headers["location"]) assert job["status"] == "failed" assert all(item["error"]["code"] == "algorithm_timeout" for item in job["items"]) assert client.get(job["download_url"]).status_code == 200 # 失败清单仍可下载。 def test_upload_validation_and_missing_jobs(client, chart, monkeypatch): assert client.post(BASE + "/batches").status_code == 422 assert client.post(BASE + "/batches", files=[("files", ("x.png", chart))] * 21).status_code == 422 monkeypatch.setattr("api.algorithm_routes.MAX_UPLOAD_BYTES", 8) assert client.post(BASE + "/batches", files=[("files", ("x.png", chart))]).status_code == 413 monkeypatch.setattr("api.algorithm_routes.MAX_UPLOAD_BYTES", len(chart)) monkeypatch.setattr("api.algorithm_routes.MAX_BATCH_BYTES", len(chart)) assert client.post(BASE + "/batches", files=[("files", ("x.png", chart))] * 2).status_code == 413 assert client.get(BASE + "/batches/missing").status_code == 404 assert client.get(BASE + "/batches/missing/download").status_code == 404 assert not list(Path(config.ALGORITHM_JOB_DIR).iterdir()) def test_restore_interrupts_incomplete_and_retains_success(tmp_path): job_id = "a" * 32 directory = tmp_path / job_id result_dir = directory / "items/0001" result_dir.mkdir(parents=True) (result_dir / "result.json").write_text('{"ok":true}') (directory / "inputs").mkdir() (directory / "inputs/0002.png").write_bytes(b"unfinished") job = {"id": job_id, "status": "running", "detailed": False, "items": [ {"id": "0001", "status": "completed", "filename": "a.png"}, {"id": "0002", "status": "running", "filename": "b.png"}, {"id": "0003", "status": "queued", "filename": "c.png"}]} (directory / "job.json").write_text(json.dumps(job)) manager = BatchManager(tmp_path) try: recovered = manager.status(job_id) assert recovered["status"] == "interrupted" assert recovered["progress"]["interrupted"] == 2 assert manager.result(job_id, "0001") == {"ok": True} with manager.open_archive(job_id) as archive: assert zipfile.is_zipfile(archive) assert not (directory / "inputs").exists() finally: manager.close() manager = BatchManager(tmp_path) try: assert manager.status(job_id)["status"] == "interrupted" # 过期后删除持久化结果，查询返回 404。 manager._jobs[job_id]["expires_at"] = (datetime.now(timezone.utc) - timedelta(seconds=1)).isoformat() manager.cleanup() assert not directory.exists() with pytest.raises(service.AlgorithmError) as exc: manager.status(job_id) assert exc.value.status_code == 404 finally: manager.close() def test_shutdown_finishes_active_and_interrupts_queued(tmp_path, chart, monkeypatch): entered, release = threading.Event(), threading.Event() def execute(image_path, result_dir, **kwargs): entered.set() assert release.wait(10) result_dir.mkdir(parents=True) (result_dir / "result.json").write_text("{}") monkeypatch.setattr(service, "execute_image", execute) manager = BatchManager(tmp_path, parallelism=1) job = manager.submit([("a.png", chart), ("b.png", chart)]) assert entered.wait(5) with ThreadPoolExecutor(1) as pool: closing = pool.submit(manager.close) try: wait_until(lambda: manager._closed) finally: release.set() closing.result(timeout=5) saved = json.loads((tmp_path / job["id"] / "job.json").read_text()) assert saved["status"] == "interrupted" assert [item["status"] for item in saved["items"]] == ["completed", "interrupted"]