vent_agent_2.0/tests/test_algorithm_batches.py

"""批量真实计算、并行上限、失败隔离、恢复和保留期限。"""
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"]