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()