X-Financial/server/tests/commercial_runtime_testkit.py

from __future__ import annotations

import uuid
from datetime import datetime, timedelta
from decimal import Decimal
from typing import Any

from sqlalchemy.orm import Session

from app.models.agent_run import AgentRun, AgentToolCall
from app.schemas.commercial import (
    CommercialEntitlementUpsert,
    CommercialPlanCreate,
    CommercialSubscriptionCreate,
)
from app.services.commercial_admin import CommercialAdminService


def seed_meter(
    db: Session,
    tenant_id: str,
    now: datetime,
    *,
    basis: str,
    tool_type: str = "llm",
    tool_name: str = "chat.completions",
    hard_limit: Decimal = Decimal("1000"),
    internal_cost: dict[str, Any] | None = None,
    subscription_status: str = "active",
    preflight_quantity: Decimal | None = None,
):
    admin = CommercialAdminService(db)
    plan = admin.create_plan(
        tenant_id,
        CommercialPlanCreate(
            plan_code="runtime",
            name="运行时计量测试套餐",
            pricing_model="hybrid",
            billing_interval="monthly",
            currency="CNY",
            base_fee=Decimal("100"),
            included_seats=10,
            effective_from=now - timedelta(days=30),
        ),
        actor_id="platform-admin",
    )
    admin.activate_plan(tenant_id, plan.id, expected_version=plan.version)
    subscription = admin.create_subscription(
        tenant_id,
        CommercialSubscriptionCreate(
            subscription_key=f"{tenant_id}-runtime",
            plan_id=plan.id,
            status=subscription_status,
            starts_at=now - timedelta(days=10),
            current_period_start=now - timedelta(days=1),
            current_period_end=now + timedelta(days=29),
            seats=5,
        ),
        actor_id="platform-admin",
    )
    runtime_meter: dict[str, Any] = {
        "enabled": True,
        "tool_type": tool_type,
        "tool_name": tool_name,
        "quantity_basis": basis,
    }
    if preflight_quantity is not None:
        runtime_meter["preflight_quantity"] = str(preflight_quantity)
    if internal_cost is not None:
        runtime_meter["internal_cost"] = internal_cost
    entitlement = admin.upsert_entitlement(
        tenant_id,
        CommercialEntitlementUpsert(
            subscription_id=subscription.id,
            entitlement_key=f"runtime_{basis}",
            metric_key=f"runtime_{basis}",
            entitlement_type="metered",
            unit="call" if basis == "call" else basis,
            included_quantity=hard_limit,
            hard_limit_quantity=hard_limit,
            reset_interval="monthly",
            overage_policy="block",
            status="active",
            effective_from=now - timedelta(days=10),
            config_json={"runtime_meter": runtime_meter},
        ),
    )
    db.flush()
    return subscription, entitlement


def seed_tool_call(
    db: Session,
    occurred_at: datetime,
    *,
    route_json: dict[str, Any],
    request_json: dict[str, Any] | None = None,
    response_json: dict[str, Any] | None = None,
    duration_ms: int = 0,
    tool_type: str = "llm",
    tool_name: str = "chat.completions",
    user_id: str = "sensitive-user-id",
    run: AgentRun | None = None,
    status: str = "succeeded",
) -> tuple[AgentRun, AgentToolCall]:
    if run is None:
        run = AgentRun(
            run_id=f"run-{uuid.uuid4().hex}",
            agent="user_agent",
            source="chat",
            user_id=user_id,
            route_json=route_json,
            permission_level="write",
            status="succeeded",
            started_at=occurred_at,
        )
        db.add(run)
    tool_call = AgentToolCall(
        id=str(uuid.uuid4()),
        run_id=run.run_id,
        tool_type=tool_type,
        tool_name=tool_name,
        request_json=request_json or {},
        response_json=response_json or {},
        status=status,
        duration_ms=duration_ms,
        created_at=occurred_at,
    )
    db.add(tool_call)
    db.flush()
    return run, tool_call


def seed_run(
    db: Session,
    started_at: datetime,
    *,
    route_json: dict[str, Any],
    status: str = "running",
) -> AgentRun:
    run = AgentRun(
        run_id=f"run-{uuid.uuid4().hex}",
        agent="user_agent",
        source="chat",
        user_id="runtime-user",
        route_json=route_json,
        permission_level="write",
        status=status,
        started_at=started_at,
    )
    db.add(run)
    db.commit()
    return run