patbond-ai/tests/providers/openai/test_tasks.py

"""同步→异步任务注册表的测试。

这一层是 gpt-image-2 与 Seedance 的根本差异所在（同步接口没有任务 ID），也是唯一
让本服务**持有状态**的地方。它写错的表现是丢任务或内存泄漏，两者都要到生产才显形，
所以逐条钉住。
"""

from __future__ import annotations

import threading
import time

import pytest

from app.domain.generation import (
    Failure,
    GenerationResult,
    Retry,
    TaskState,
    UserFacing,
)
from app.providers.openai.client import OpenAiError
from app.providers.openai.tasks import LOCAL_TASK_PREFIX, SyncTaskRegistry


def succeeded(url: str = "data:image/png;base64,AAA") -> GenerationResult:
    return GenerationResult(
        provider_request_id="", state=TaskState.SUCCEEDED, progress=100, asset_url=url
    )


def wait_until_done(registry: SyncTaskRegistry, task_id: str, timeout: float = 3.0):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        result = registry.fetch(task_id)
        if result is not None and result.state.is_terminal:
            return result
        time.sleep(0.01)
    raise AssertionError(f"任务 {task_id} 在 {timeout}s 内未结束")


class TestSubmitReturnsImmediately:
    def test_submit_does_not_block_on_the_slow_call(self):
        """这是整层存在的理由：同步调用必须在后台跑，提交立刻返回。"""
        registry = SyncTaskRegistry(max_workers=2)
        started = threading.Event()
        release = threading.Event()

        def slow():
            started.set()
            release.wait(timeout=3)
            return succeeded()

        begin = time.monotonic()
        task_id = registry.submit(slow)
        elapsed = time.monotonic() - begin

        assert started.wait(timeout=2)
        # 提交本身几乎不耗时，尽管那个调用还在跑。
        assert elapsed < 0.5
        assert registry.fetch(task_id).state is TaskState.RUNNING
        release.set()
        wait_until_done(registry, task_id)
        registry.shutdown()

    def test_task_id_is_prefixed_so_logs_distinguish_it_from_a_vendor_id(self):
        registry = SyncTaskRegistry()
        task_id = registry.submit(succeeded)
        assert task_id.startswith(LOCAL_TASK_PREFIX)
        registry.shutdown()

    def test_ids_are_unique(self):
        registry = SyncTaskRegistry()
        ids = {registry.submit(succeeded) for _ in range(20)}
        assert len(ids) == 20
        registry.shutdown()


class TestInFlightAndCompletedStates:
    def test_running_reports_zero_progress(self):
        # 同步接口不提供进度，只能报 0。契约本就不承诺平滑进度（ADR-023 §10）。
        registry = SyncTaskRegistry()
        release = threading.Event()
        task_id = registry.submit(lambda: (release.wait(timeout=3), succeeded())[1])

        result = registry.fetch(task_id)
        assert result.state is TaskState.RUNNING
        assert result.progress == 0
        release.set()
        registry.shutdown()

    def test_completed_result_is_returned_with_the_real_task_id(self):
        """工作函数拿不到自己的任务 ID，故由注册表在 fetch 时填入。"""
        registry = SyncTaskRegistry()
        task_id = registry.submit(succeeded)

        result = wait_until_done(registry, task_id)
        assert result.state is TaskState.SUCCEEDED
        assert result.provider_request_id == task_id
        assert result.asset_url is not None
        registry.shutdown()

    def test_result_can_be_fetched_more_than_once(self):
        # Java 侧可能重复轮询同一个已完成任务（比如一次响应丢了）。
        registry = SyncTaskRegistry()
        task_id = registry.submit(succeeded)
        wait_until_done(registry, task_id)

        assert registry.fetch(task_id).state is TaskState.SUCCEEDED
        assert registry.fetch(task_id).state is TaskState.SUCCEEDED
        registry.shutdown()


class TestUnknownIdIsIndistinguishableOnPurpose:
    def test_never_existed_returns_none(self):
        """从未存在过与「进程重启后丢失」无法区分——不假装区分。

        两种情形对调用方的处置相同：停止轮询、任务判失败。
        """
        assert SyncTaskRegistry().fetch("local-nonexistent") is None


class TestFailurePropagation:
    def test_a_classified_provider_error_keeps_its_classification(self):
        registry = SyncTaskRegistry()
        failure = Failure(
            Retry.QUOTA, UserFacing.QUOTA_EXHAUSTED, "insufficient_quota", "no credit"
        )

        def failing():
            raise OpenAiError(failure, 429)

        task_id = registry.submit(failing)
        result = wait_until_done(registry, task_id)

        assert result.state is TaskState.FAILED
        assert result.failure.retry is Retry.QUOTA
        assert result.failure.provider_code == "insufficient_quota"
        registry.shutdown()

    def test_an_unclassified_exception_becomes_a_retryable_failure_naming_its_type(
        self,
    ):
        """未分类异常不该让任务卡住，也不该丢掉线索。

        把异常类型名放进 provider_code——否则排查时只有一句空洞的错误消息。
        """
        registry = SyncTaskRegistry()

        def boom():
            raise ZeroDivisionError("unexpected")

        task_id = registry.submit(boom)
        result = wait_until_done(registry, task_id)

        assert result.state is TaskState.FAILED
        assert result.failure.retry is Retry.BACKOFF
        assert result.failure.provider_code == "ZeroDivisionError"
        registry.shutdown()


class TestCapacityIsBounded:
    def test_submitting_beyond_the_cap_is_rejected_rather_than_evicting(self):
        """满了要拒绝新提交，**不能挤掉在途任务**。

        挤掉一个在途任务会让一个可能正在计费的调用变成无人认领的开销。
        """
        registry = SyncTaskRegistry(max_workers=1, max_entries=2)
        release = threading.Event()

        def blocked():
            release.wait(timeout=3)
            return succeeded()

        registry.submit(blocked)
        registry.submit(blocked)

        with pytest.raises(RuntimeError, match="上限"):
            registry.submit(blocked)

        release.set()
        registry.shutdown()

    def test_concurrency_is_capped_by_max_workers(self):
        """并发上限同时也是「同时有多少钱在花」的上界。"""
        registry = SyncTaskRegistry(max_workers=2, max_entries=10)
        concurrent = []
        lock = threading.Lock()
        release = threading.Event()

        def work():
            with lock:
                concurrent.append(1)
                peak = len(concurrent)
            release.wait(timeout=3)
            with lock:
                concurrent.pop()
            return succeeded(f"data:image/png;base64,{peak}")

        for _ in range(5):
            registry.submit(work)
        time.sleep(0.2)

        with lock:
            assert len(concurrent) <= 2

        release.set()
        registry.shutdown()


class TestExpiryOnlyTouchesCompletedEntries:
    def test_completed_entries_expire(self):
        registry = SyncTaskRegistry(completed_ttl_sec=0.05, max_entries=10)
        task_id = registry.submit(succeeded)
        wait_until_done(registry, task_id)
        time.sleep(0.1)

        # 清理发生在下一次 submit 时（不额外起一个定时线程——本服务的规模用不上）。
        registry.submit(succeeded)
        assert registry.fetch(task_id) is None
        registry.shutdown()

    def test_in_flight_entries_are_never_expired(self):
        """在途任务无论多久都不清——它可能正在花钱。"""
        registry = SyncTaskRegistry(completed_ttl_sec=0.01, max_entries=10)
        release = threading.Event()
        task_id = registry.submit(lambda: (release.wait(timeout=3), succeeded())[1])
        time.sleep(0.1)

        registry.submit(succeeded)
        # 仍在注册表里。
        assert registry.fetch(task_id) is not None
        assert registry.fetch(task_id).state is TaskState.RUNNING

        release.set()
        registry.shutdown()

    def test_ttl_starts_at_completion_not_at_creation(self):
        """否则一个耗时接近 TTL 的任务会在刚完成时就被清掉——结果还没被取走。"""
        registry = SyncTaskRegistry(completed_ttl_sec=5.0, max_entries=10)
        release = threading.Event()
        task_id = registry.submit(lambda: (release.wait(timeout=3), succeeded())[1])
        time.sleep(0.15)
        release.set()
        wait_until_done(registry, task_id)

        registry.submit(succeeded)
        assert registry.fetch(task_id).state is TaskState.SUCCEEDED
        registry.shutdown()