X-Financial/server/tests/test_agent_asset_release_scheduler.py

from __future__ import annotations

import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from app.core.agent_enums import AgentAssetDomain, AgentAssetStatus, AgentAssetType
from app.db.base import Base
from app.models.agent_asset import AgentAsset, AgentAssetTestRun
from app.services.agent_asset_release_scheduler import AgentAssetReleaseScheduler
from app.services.agent_asset_release_telemetry import (
    AgentAssetReleaseTelemetryService,
    ReleaseObservationInput,
)


def _asset(asset_id: str, tenant_id: str | None) -> AgentAsset:
    previous_config: dict[str, object] = {
        "detail_mode": "json_risk",
        "enabled": True,
        "stable_marker": f"stable:{asset_id}",
    }
    if tenant_id is not None:
        previous_config["tenant_id"] = tenant_id
    return AgentAsset(
        id=asset_id,
        tenant_id=tenant_id or "platform",
        scope="tenant" if tenant_id is not None else "platform",
        asset_type=AgentAssetType.RULE.value,
        code=f"risk.{asset_id}",
        name="周期发布监控规则",
        description="",
        domain=AgentAssetDomain.EXPENSE.value,
        scenario_json=["travel"],
        owner="finance",
        status=AgentAssetStatus.ACTIVE.value,
        current_version="v1",
        published_version="v1",
        working_version="v2",
        config_json={
            **previous_config,
            "release_guard": {
                "release_id": f"release-{asset_id}",
                "stage": "shadow",
                "candidate_version": "v2",
                "previous_version": "v1",
                "previous_config": previous_config,
                "policy": {
                    "shadow_min_samples": 1,
                    "canary_min_samples": 1,
                    "max_error_rate": 1.0,
                    "min_precision": 0.9,
                    "max_precision_drop": 1.0,
                    "canary_traffic_percent": 10,
                },
            },
        },
    )


def _false_positive_sample(db: Session, asset: AgentAsset, tenant_id: str) -> None:
    telemetry = AgentAssetReleaseTelemetryService(db)
    observation = telemetry.record_observation(
        ReleaseObservationInput(
            tenant_id=tenant_id,
            asset_id=asset.id,
            release_id=f"release-{asset.id}",
            stage="shadow",
            version="v2",
            rule_code=asset.code,
            source_key=f"claim-{asset.id}",
            candidate_hit=True,
            baseline_hit=True,
        )
    )
    telemetry.record_review_label(
        tenant_id=tenant_id,
        observation_id=observation.id,
        label="false_positive",
        request_id=f"review-{asset.id}",
        actor_id="trusted-reviewer",
    )


def test_scheduler_groups_tenants_rolls_back_and_skips_global_assets() -> None:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
    with factory() as db:
        tenant_asset = _asset("tenant-release-scheduler", "tenant-a")
        global_asset = _asset("global-release-scheduler", None)
        db.add_all([tenant_asset, global_asset])
        db.commit()
        _false_positive_sample(db, tenant_asset, "tenant-a")
        with pytest.raises(LookupError, match="Agent asset not found"):
            _false_positive_sample(db, global_asset, "tenant-a")
        db.commit()

    result = AgentAssetReleaseScheduler(session_factory=factory)._run_once()

    with factory() as db:
        tenant_asset = db.get(AgentAsset, "tenant-release-scheduler")
        global_asset = db.get(AgentAsset, "global-release-scheduler")
        assert tenant_asset is not None
        assert global_asset is not None
        assert tenant_asset.config_json["release_guard"]["stage"] == "rolled_back"
        assert global_asset.config_json["release_guard"]["stage"] == "shadow"
    engine.dispose()

    assert result == {
        "tenants": 1,
        "scanned": 1,
        "evaluated": 1,
        "collecting": 0,
        "rolled_back": 1,
        "errors": 0,
        "global_skipped": 1,
    }


def test_scheduler_cursor_rotates_beyond_fixed_batch_limit() -> None:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
    with factory() as db:
        first = _asset("cursor-a", "tenant-a")
        second = _asset("cursor-b", "tenant-a")
        db.add_all([first, second])
        db.commit()
        for asset in (first, second):
            telemetry = AgentAssetReleaseTelemetryService(db)
            observation = telemetry.record_observation(
                ReleaseObservationInput(
                    tenant_id="tenant-a",
                    asset_id=asset.id,
                    release_id=f"release-{asset.id}",
                    stage="shadow",
                    version="v2",
                    rule_code=asset.code,
                    source_key=f"claim-{asset.id}",
                    candidate_hit=True,
                    baseline_hit=True,
                )
            )
            telemetry.record_review_label(
                tenant_id="tenant-a",
                observation_id=observation.id,
                label="confirmed",
                request_id=f"review-{asset.id}",
                actor_id="trusted-reviewer",
            )
        db.commit()

    scheduler = AgentAssetReleaseScheduler(session_factory=factory)
    scheduler._batch_size = 1
    scheduler._run_once()
    scheduler._run_once()

    with factory() as db:
        evaluated = set(db.scalars(select(AgentAssetTestRun.asset_id)).all())
    engine.dispose()

    assert evaluated == {"cursor-a", "cursor-b"}