SupplyAgent/tests/supply_infra/pipeline/test_aigc_outbox.py

from __future__ import annotations

from contextlib import contextmanager
from unittest.mock import MagicMock

from sqlalchemy import create_engine, select
from sqlalchemy.orm import sessionmaker

from supply_infra.db.base import Base
from supply_infra.db.models.pipeline_outbox import PipelineOutbox
from supply_infra.db.models.pipeline_run import PipelineRun
from supply_infra.db.models.pipeline_step_run import PipelineStepRun
from supply_infra.pipeline import aigc_outbox
from supply_infra.pipeline.contracts import StepContext
from supply_infra.pipeline.enums import OutboxStatus


def _payload() -> dict:
    return {
        "biz_dt": "20260730",
        "source_run_id": None,
        "plan_label": "plan-a",
        "produce_plan_id": "produce-1",
        "publish_plan_id": "publish-1",
        "aweme_ids": ["aweme-1"],
        "candidate_ids": [101],
        "plan_name": "【SupplyAgent】plan-a-20260730-stable",
        "batch_key": "a" * 64,
    }


def _context(*, dry_run: bool = False) -> StepContext:
    return StepContext(
        run_id="run-1",
        step_run_id="step-1",
        step_key="aigc_write_record",
        biz_dt="20260730",
        date_snapshot={},
        config_snapshot={},
        input_snapshot={},
        dry_run=dry_run,
    )


def _database(monkeypatch):
    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(
        engine,
        tables=[
            PipelineRun.__table__,
            PipelineStepRun.__table__,
            PipelineOutbox.__table__,
        ],
    )
    session_factory = sessionmaker(bind=engine, expire_on_commit=False)

    @contextmanager
    def get_test_session():
        with session_factory() as session:
            try:
                yield session
                session.commit()
            except Exception:
                session.rollback()
                raise

    monkeypatch.setattr(aigc_outbox, "get_session", get_test_session)
    return session_factory


def _prepared() -> dict:
    return {
        "biz_dt": "20260730",
        "run_id": None,
        "plan_count": 1,
        "candidate_count": 1,
        "decision_buckets": ["primary"],
        "distribution": [],
        "batches": [_payload()],
    }


def test_event_is_committed_before_external_call_and_replay_is_idempotent(
    monkeypatch,
) -> None:
    session_factory = _database(monkeypatch)
    monkeypatch.setattr(aigc_outbox, "prepare_publish_batches", lambda **_: _prepared())
    service = MagicMock()
    service.mark_candidates_aigc_plans.return_value = 1
    monkeypatch.setattr(aigc_outbox, "get_video_discovery_service", lambda: service)

    calls = {"create": 0}

    class Client:
        def create_video_crawler_plan(self, _aweme_ids, *, plan_name):
            calls["create"] += 1
            with session_factory() as session:
                row = session.scalar(select(PipelineOutbox))
                assert row is not None
                assert row.status == OutboxStatus.SENDING.value
                assert row.payload_json["plan_name"] == plan_name
            return {
                "success": True,
                "crawler_plan_id": "crawler-1",
                "crawler_plan_name": plan_name,
            }

        def bind_crawler_to_produce_plan(self, *_args, **_kwargs):
            return {"success": True}

    monkeypatch.setattr(aigc_outbox, "AigcClient", Client)

    first = aigc_outbox.execute_aigc_outbox(_context())
    second = aigc_outbox.execute_aigc_outbox(_context())

    assert first["success"] is True
    assert second["success"] is True
    assert calls["create"] == 1
    with session_factory() as session:
        row = session.scalar(select(PipelineOutbox))
        assert row is not None
        assert row.status == OutboxStatus.SUCCEEDED.value
        assert row.external_id == "crawler-1"
        assert row.attempt_count == 1


def test_dry_run_records_batches_without_calling_aigc(monkeypatch) -> None:
    session_factory = _database(monkeypatch)
    monkeypatch.setattr(
        aigc_outbox,
        "prepare_publish_batches",
        lambda **_: _prepared(),
    )
    client = MagicMock(side_effect=AssertionError("AIGC must not be called"))
    monkeypatch.setattr(aigc_outbox, "AigcClient", client)

    result = aigc_outbox.execute_aigc_outbox(_context(dry_run=True))

    assert result["success"] is True
    assert result["dry_run"] is True
    assert result["effect_recorded"] is True
    assert result["external_request_made"] is False
    client.assert_not_called()
    with session_factory() as session:
        row = session.scalar(select(PipelineOutbox))
        assert row is not None
        assert row.dry_run is True
        assert row.status == OutboxStatus.SUCCEEDED.value
        assert row.attempt_count == 0
        assert row.external_id is None
        assert row.response_json["external_request_made"] is False


def test_timeout_becomes_ambiguous_and_is_not_automatically_replayed(
    monkeypatch,
) -> None:
    session_factory = _database(monkeypatch)
    service = MagicMock()
    monkeypatch.setattr(aigc_outbox, "get_video_discovery_service", lambda: service)
    event_id = aigc_outbox._enqueue_batches(_context(), [_payload()])[0]

    class TimeoutClient:
        calls = 0

        def create_video_crawler_plan(self, *_args, **_kwargs):
            self.calls += 1
            raise TimeoutError("remote timeout")

    client = TimeoutClient()
    first = aigc_outbox.dispatch_aigc_outbox_event(event_id, client=client)
    second = aigc_outbox.dispatch_aigc_outbox_event(event_id, client=client)

    assert first["status"] == OutboxStatus.AMBIGUOUS.value
    assert second["status"] == OutboxStatus.AMBIGUOUS.value
    assert client.calls == 1
    service.mark_candidates_aigc_plans.assert_not_called()
    with session_factory() as session:
        row = session.get(PipelineOutbox, event_id)
        assert row is not None
        assert row.status == OutboxStatus.AMBIGUOUS.value


def test_bind_retry_reuses_persisted_external_plan(monkeypatch) -> None:
    session_factory = _database(monkeypatch)
    service = MagicMock()
    service.mark_candidates_aigc_plans.return_value = 1
    monkeypatch.setattr(aigc_outbox, "get_video_discovery_service", lambda: service)
    event_id = aigc_outbox._enqueue_batches(_context(), [_payload()])[0]

    first_client = MagicMock()
    first_client.create_video_crawler_plan.return_value = {
        "success": True,
        "crawler_plan_id": "crawler-1",
        "crawler_plan_name": "stable-name",
    }
    first_client.bind_crawler_to_produce_plan.side_effect = TimeoutError("bind timeout")

    first = aigc_outbox.dispatch_aigc_outbox_event(event_id, client=first_client)
    assert first["status"] == OutboxStatus.RETRYABLE_FAILED.value

    retry_client = MagicMock()
    retry_client.bind_crawler_to_produce_plan.return_value = {"success": True}
    second = aigc_outbox.dispatch_aigc_outbox_event(event_id, client=retry_client)

    assert second["success"] is True
    retry_client.create_video_crawler_plan.assert_not_called()
    retry_client.bind_crawler_to_produce_plan.assert_called_once()
    with session_factory() as session:
        row = session.get(PipelineOutbox, event_id)
        assert row is not None
        assert row.status == OutboxStatus.SUCCEEDED.value
        assert row.external_id == "crawler-1"
        assert row.attempt_count == 2