X-Financial/server/tests/commercial_migration_assertions.py

from __future__ import annotations

from typing import Any

import pytest
from sqlalchemy import inspect, text
from sqlalchemy.engine import Engine
from sqlalchemy.exc import DBAPIError, IntegrityError

from app.models.commercial import (
    CommercialCostEvent,
    CommercialEntitlement,
    TenantCommercialPlan,
    TenantSubscription,
    UsageMeterEvent,
)
from app.models.commercial_billing import CommercialAdminEvent, CommercialBillingPeriod
from app.models.commercial_runtime import CommercialRuntimeReservation

COMMERCIAL_MODELS = (
    TenantCommercialPlan,
    TenantSubscription,
    CommercialEntitlement,
    UsageMeterEvent,
    CommercialCostEvent,
    CommercialRuntimeReservation,
    CommercialBillingPeriod,
    CommercialAdminEvent,
)


def _assert_commercial_head_schema(engine: Engine) -> None:
    inspector = inspect(engine)
    for model in COMMERCIAL_MODELS:
        table = model.__table__
        live_columns = {
            str(column["name"]): bool(column["nullable"])
            for column in inspector.get_columns(table.name, schema="public")
        }
        declared_columns = {column.name: bool(column.nullable) for column in table.columns}
        assert live_columns == declared_columns

        live_constraint_names = {
            str(item["name"])
            for loader in (
                inspector.get_unique_constraints,
                inspector.get_check_constraints,
                inspector.get_foreign_keys,
            )
            for item in loader(table.name, schema="public")
            if item.get("name")
        }
        declared_constraint_names = {
            str(constraint.name) for constraint in table.constraints if constraint.name is not None
        }
        assert live_constraint_names == declared_constraint_names

        live_index_names = {
            str(item["name"])
            for item in inspector.get_indexes(table.name, schema="public")
            if not item.get("duplicates_constraint")
        }
        declared_index_names = {str(index.name) for index in table.indexes}
        assert live_index_names == declared_index_names

    savings_tables = {
        "profile_baseline_snapshots",
        "savings_opportunities",
        "savings_realizations",
        "savings_evidence_links",
        "savings_events",
    }
    for table_name in ("usage_meter_events", "commercial_cost_events"):
        target_tables = {
            str(foreign_key["referred_table"])
            for foreign_key in inspector.get_foreign_keys(table_name, schema="public")
        }
        assert target_tables.isdisjoint(savings_tables)

    with engine.connect() as connection:
        active_plan_index = str(
            connection.scalar(
                text(
                    "SELECT indexdef FROM pg_indexes WHERE schemaname = 'public' "
                    "AND tablename = 'tenant_commercial_plans' "
                    "AND indexname = 'uq_tenant_commercial_plans_active_code'"
                )
            )
            or ""
        ).lower()
        current_subscription_index = str(
            connection.scalar(
                text(
                    "SELECT indexdef FROM pg_indexes WHERE schemaname = 'public' "
                    "AND tablename = 'tenant_subscriptions' "
                    "AND indexname = 'uq_tenant_subscriptions_current'"
                )
            )
            or ""
        ).lower()
        trigger_counts = {
            table_name: int(
                connection.scalar(
                    text(
                        "SELECT COUNT(*) FROM pg_trigger trigger "
                        "JOIN pg_class relation ON relation.oid = trigger.tgrelid "
                        "WHERE relation.relname = :table_name "
                        "AND trigger.tgname = :trigger_name "
                        "AND NOT trigger.tgisinternal"
                    ),
                    {
                        "table_name": table_name,
                        "trigger_name": f"trg_{table_name}_append_only",
                    },
                )
                or 0
            )
            for table_name in (
                "usage_meter_events",
                "commercial_cost_events",
                "commercial_billing_periods",
                "commercial_admin_events",
            )
        }
        period_overlap_trigger_count = int(
            connection.scalar(
                text(
                    "SELECT COUNT(*) FROM pg_trigger trigger "
                    "JOIN pg_class relation ON relation.oid = trigger.tgrelid "
                    "WHERE relation.relname = 'commercial_billing_periods' "
                    "AND trigger.tgname = "
                    "'trg_commercial_billing_periods_no_overlap' "
                    "AND NOT trigger.tgisinternal"
                )
            )
            or 0
        )
    assert "unique index" in active_plan_index and "status" in active_plan_index
    assert "active" in active_plan_index
    assert "unique index" in current_subscription_index
    assert all(status in current_subscription_index for status in ("active", "suspended"))
    assert trigger_counts == {
        "usage_meter_events": 1,
        "commercial_cost_events": 1,
        "commercial_billing_periods": 1,
        "commercial_admin_events": 1,
    }
    assert period_overlap_trigger_count == 1


def _assert_commercial_runtime_invariants(engine: Engine) -> None:
    def execute_rejected(
        connection: Any,
        statement: Any,
        parameters: dict[str, Any],
        error_type: type[DBAPIError] = IntegrityError,
    ) -> None:
        savepoint = connection.begin_nested()
        try:
            with pytest.raises(error_type):
                connection.execute(statement, parameters)
        finally:
            if savepoint.is_active:
                savepoint.rollback()

    subscription_insert = text(
        """
        INSERT INTO tenant_subscriptions (
            id, tenant_id, subscription_key, plan_id, status, starts_at,
            current_period_start, current_period_end, seats, base_fee_snapshot,
            currency, billing_interval, created_by
        ) VALUES (
            :id, :tenant_id, :subscription_key, :plan_id, 'active', now(),
            now(), now() + interval '1 month', 10, 1000, 'CNY', 'monthly', 'probe'
        )
        """
    )
    usage_insert = text(
        """
        INSERT INTO usage_meter_events (
            id, tenant_id, subscription_id, entitlement_id, billing_period_id,
            event_type, metric_key, quantity, unit, period_key, quota_period_key,
            occurred_at, source_system,
            idempotency_key, request_fingerprint, actor_type, actor_id
        ) VALUES (
            :id, :tenant_id, :subscription_id, :entitlement_id, 'period-a', 'usage',
            'ai.review', 1, 'request', 'bp-probe', '2026-07', now(), 'runtime-probe',
            :idempotency_key, :request_fingerprint, 'system', 'migration-probe'
        )
        """
    )
    cost_insert = text(
        """
        INSERT INTO commercial_cost_events (
            id, tenant_id, subscription_id, billing_period_id, usage_event_id, event_type,
            cost_category, quantity, unit, unit_cost, cost_amount,
            original_currency, reporting_amount, reporting_currency, fx_rate,
            allocation_key, occurred_at, source_system, idempotency_key,
            request_fingerprint
        ) VALUES (
            :id, :tenant_id, :subscription_id, 'period-a', :usage_event_id, 'incurred',
            'ai_inference', 10, 'token', 0.01, 0.10,
            'CNY', 0.10, 'CNY', 1, 'ai.review', now(), 'runtime-probe',
            :idempotency_key, :request_fingerprint
        )
        """
    )

    with engine.connect() as connection:
        transaction = connection.begin()
        try:
            connection.execute(
                text(
                    """
                    INSERT INTO tenant_commercial_plans (
                        id, tenant_id, plan_code, name, pricing_model,
                        billing_interval, currency, base_fee, included_seats,
                        status, effective_from, created_by
                    ) VALUES (
                        'plan-a', 'tenant-a', 'enterprise', '企业版', 'hybrid',
                        'monthly', 'CNY', 1000, 10, 'active', now(), 'probe'
                    )
                    """
                )
            )
            connection.execute(
                subscription_insert,
                {
                    "id": "subscription-a",
                    "tenant_id": "tenant-a",
                    "subscription_key": "subscription-a",
                    "plan_id": "plan-a",
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO commercial_billing_periods (
                        id, tenant_id, subscription_id, plan_id, period_sequence,
                        period_key, status, period_start, period_end,
                        subscription_status_snapshot, plan_code_snapshot,
                        plan_version_snapshot, pricing_model_snapshot,
                        billing_interval, currency, base_fee_snapshot,
                        seats_snapshot, source, idempotency_key, created_by
                    ) VALUES (
                        'period-a', 'tenant-a', 'subscription-a', 'plan-a', 1,
                        'bp-probe', 'issued', now(), now() + interval '1 month',
                        'active', 'enterprise', 1, 'hybrid', 'monthly', 'CNY',
                        1000, 10, 'subscription_created', 'period-request-a', 'probe'
                    )
                    """
                )
            )
            execute_rejected(
                connection,
                text(
                    """
                    INSERT INTO commercial_billing_periods (
                        id, tenant_id, subscription_id, plan_id, period_sequence,
                        period_key, status, period_start, period_end,
                        subscription_status_snapshot, plan_code_snapshot,
                        plan_version_snapshot, pricing_model_snapshot,
                        billing_interval, currency, base_fee_snapshot,
                        seats_snapshot, source, idempotency_key, created_by
                    ) VALUES (
                        'period-overlap', 'tenant-a', 'subscription-a', 'plan-a', 2,
                        'bp-overlap', 'issued', now() + interval '15 days',
                        now() + interval '45 days', 'active', 'enterprise', 1,
                        'hybrid', 'monthly', 'CNY', 1000, 10, 'auto_renew',
                        'period-request-overlap', 'probe'
                    )
                    """
                ),
                {},
                DBAPIError,
            )
            connection.execute(
                text(
                    """
                    INSERT INTO commercial_admin_events (
                        id, tenant_id, actor_type, actor_id, request_id, reason,
                        action, resource_type, resource_id, resource_version,
                        before_json, after_json
                    ) VALUES (
                        'admin-event-a', 'tenant-a', 'user', 'probe',
                        'admin-request-a', '迁移不变量探针', 'billing_period_created',
                        'billing_period', 'period-a', 1, '{}'::json, '{}'::json
                    )
                    """
                )
            )
            execute_rejected(
                connection,
                subscription_insert,
                {
                    "id": "cross-tenant-subscription",
                    "tenant_id": "tenant-b",
                    "subscription_key": "cross-tenant-subscription",
                    "plan_id": "plan-a",
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO commercial_entitlements (
                        id, tenant_id, subscription_id, entitlement_key,
                        metric_key, entitlement_type, unit, included_quantity,
                        hard_limit_quantity, reset_interval, overage_policy,
                        status, effective_from
                    ) VALUES (
                        'entitlement-a', 'tenant-a', 'subscription-a', 'ai-review',
                        'ai.review', 'metered', 'request', 100, 120, 'monthly',
                        'block', 'active', now()
                    )
                    """
                )
            )
            connection.execute(
                text(
                    """
                    INSERT INTO commercial_runtime_reservations (
                        id, tenant_id, subscription_id, entitlement_id,
                        billing_period_id, run_id,
                        tool_call_id, tool_type, tool_name, quantity_basis,
                        reserved_quantity, period_key, quota_period_key, status,
                        request_fingerprint,
                        meter_config_json, expires_at
                    ) VALUES (
                        'reservation-a', 'tenant-a', 'subscription-a',
                        'entitlement-a', 'period-a', 'run-a', 'tool-call-a', 'llm',
                        'chat.completions', 'call', 1, 'bp-probe', '2026-07', 'reserved',
                        'reservation-fingerprint-a', '{}'::json, now() + interval '15 minutes'
                    )
                    """
                )
            )
            execute_rejected(
                connection,
                text(
                    """
                    INSERT INTO commercial_runtime_reservations (
                        id, tenant_id, subscription_id, entitlement_id,
                        billing_period_id, run_id,
                        tool_call_id, tool_type, tool_name, quantity_basis,
                        reserved_quantity, period_key, quota_period_key, status,
                        request_fingerprint,
                        meter_config_json, expires_at
                    ) VALUES (
                        'reservation-duplicate', 'tenant-a', 'subscription-a',
                        'entitlement-a', 'period-a', 'run-b', 'tool-call-a', 'llm',
                        'chat.completions', 'call', 1, 'bp-probe', '2026-07', 'reserved',
                        'reservation-fingerprint-b', '{}'::json, now() + interval '15 minutes'
                    )
                    """
                ),
                {},
            )
            execute_rejected(
                connection,
                text(
                    """
                    INSERT INTO commercial_runtime_reservations (
                        id, tenant_id, subscription_id, entitlement_id,
                        billing_period_id, run_id,
                        tool_call_id, tool_type, tool_name, quantity_basis,
                        reserved_quantity, actual_quantity, period_key,
                        quota_period_key, status,
                        request_fingerprint, meter_config_json, expires_at, settled_at
                    ) VALUES (
                        'reservation-invalid', 'tenant-a', 'subscription-a',
                        'entitlement-a', 'period-a', 'run-c', 'tool-call-c', 'llm',
                        'chat.completions', 'call', 1, 2, 'bp-probe', '2026-07', 'committed',
                        'reservation-fingerprint-c', '{}'::json,
                        now() + interval '15 minutes', now()
                    )
                    """
                ),
                {},
            )
            connection.execute(
                usage_insert,
                {
                    "id": "usage-a",
                    "tenant_id": "tenant-a",
                    "subscription_id": "subscription-a",
                    "entitlement_id": "entitlement-a",
                    "idempotency_key": "usage-request-a",
                    "request_fingerprint": "sha256:usage-a",
                },
            )
            execute_rejected(
                connection,
                usage_insert,
                {
                    "id": "usage-duplicate",
                    "tenant_id": "tenant-a",
                    "subscription_id": "subscription-a",
                    "entitlement_id": "entitlement-a",
                    "idempotency_key": "usage-request-a",
                    "request_fingerprint": "sha256:different-payload",
                },
            )
            execute_rejected(
                connection,
                usage_insert,
                {
                    "id": "usage-cross-tenant",
                    "tenant_id": "tenant-b",
                    "subscription_id": "subscription-a",
                    "entitlement_id": "entitlement-a",
                    "idempotency_key": "usage-cross-tenant",
                    "request_fingerprint": "sha256:cross-tenant",
                },
            )
            connection.execute(
                cost_insert,
                {
                    "id": "cost-a",
                    "tenant_id": "tenant-a",
                    "subscription_id": "subscription-a",
                    "usage_event_id": "usage-a",
                    "idempotency_key": "cost-request-a",
                    "request_fingerprint": "sha256:cost-a",
                },
            )
            execute_rejected(
                connection,
                cost_insert,
                {
                    "id": "cost-duplicate",
                    "tenant_id": "tenant-a",
                    "subscription_id": "subscription-a",
                    "usage_event_id": "usage-a",
                    "idempotency_key": "cost-request-a",
                    "request_fingerprint": "sha256:different-cost",
                },
            )
            execute_rejected(
                connection,
                cost_insert,
                {
                    "id": "cost-cross-tenant",
                    "tenant_id": "tenant-b",
                    "subscription_id": "subscription-a",
                    "usage_event_id": "usage-a",
                    "idempotency_key": "cost-cross-tenant",
                    "request_fingerprint": "sha256:cross-cost",
                },
            )
            execute_rejected(
                connection,
                text("UPDATE usage_meter_events SET quantity = 2 WHERE id = 'usage-a'"),
                {},
                DBAPIError,
            )
            execute_rejected(
                connection,
                text("DELETE FROM commercial_cost_events WHERE id = 'cost-a'"),
                {},
                DBAPIError,
            )
            execute_rejected(
                connection,
                text(
                    "UPDATE commercial_billing_periods SET currency = 'USD' WHERE id = 'period-a'"
                ),
                {},
                DBAPIError,
            )
            execute_rejected(
                connection,
                text("DELETE FROM commercial_admin_events WHERE id = 'admin-event-a'"),
                {},
                DBAPIError,
            )
        finally:
            transaction.rollback()