from __future__ import annotations import threading import uuid from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime, timedelta from decimal import Decimal from savings_postgres_testkit import ( # noqa: F401 - 注册 pg_factory fixture _pg_factory_fixture, ) from sqlalchemy import func, select from sqlalchemy.orm import Session, sessionmaker from app.models.commercial import CommercialCostEvent, UsageMeterEvent from app.models.commercial_billing import CommercialBillingPeriod from app.models.commercial_runtime import CommercialRuntimeReservation from app.schemas.commercial import ( CommercialCostEventCreate, CommercialEntitlementUpsert, CommercialPlanCreate, CommercialSubscriptionCreate, UsageMeterEventCreate, ) from app.services.commercial_access_policy import CommercialConflictError from app.services.commercial_admin import CommercialAdminService from app.services.commercial_metering import CommercialMeteringService from app.services.commercial_runtime_reservations import ( CommercialRuntimeReservationService, ) from app.services.commercial_subscription_rollover import ( CommercialSubscriptionRolloverService, ) def test_concurrent_usage_replay_creates_one_event( pg_factory: sessionmaker[Session], ) -> None: tenant_id, subscription_id, entitlement_id = _seed_account(pg_factory) payload = UsageMeterEventCreate( subscription_id=subscription_id, entitlement_id=entitlement_id, quantity=Decimal("2"), occurred_at=datetime.now(UTC), source_system="postgres-runtime", idempotency_key=f"usage-{uuid.uuid4().hex}", ) ready = threading.Barrier(2) def record_once() -> bool: with pg_factory() as db: ready.wait(timeout=5) _, created = CommercialMeteringService(db).record_usage( tenant_id, payload, actor_type="system", actor_id="postgres-runtime", ) db.commit() return created with ThreadPoolExecutor(max_workers=2) as pool: outcomes = [ future.result(timeout=10) for future in (pool.submit(record_once), pool.submit(record_once)) ] assert sorted(outcomes) == [False, True] with pg_factory() as db: assert ( db.scalar( select(func.count()) .select_from(UsageMeterEvent) .where( UsageMeterEvent.tenant_id == tenant_id, UsageMeterEvent.idempotency_key == payload.idempotency_key, ) ) == 1 ) def test_concurrent_usage_cannot_bypass_hard_limit( pg_factory: sessionmaker[Session], ) -> None: tenant_id, subscription_id, entitlement_id = _seed_account(pg_factory) ready = threading.Barrier(2) def consume_once(suffix: str) -> str: payload = UsageMeterEventCreate( subscription_id=subscription_id, entitlement_id=entitlement_id, quantity=Decimal("4"), occurred_at=datetime.now(UTC), source_system="postgres-runtime", idempotency_key=f"quota-{suffix}-{uuid.uuid4().hex}", ) with pg_factory() as db: ready.wait(timeout=5) try: CommercialMeteringService(db).record_usage( tenant_id, payload, actor_type="system", actor_id="postgres-runtime", ) db.commit() return "created" except CommercialConflictError: db.rollback() return "blocked" with ThreadPoolExecutor(max_workers=2) as pool: outcomes = [ future.result(timeout=10) for future in (pool.submit(consume_once, "a"), pool.submit(consume_once, "b")) ] assert sorted(outcomes) == ["blocked", "created"] with pg_factory() as db: used = db.scalar( select(func.sum(UsageMeterEvent.quantity)).where( UsageMeterEvent.tenant_id == tenant_id, UsageMeterEvent.entitlement_id == entitlement_id, ) ) assert Decimal(used or 0) == Decimal("4.000000") def test_concurrent_cost_reversal_has_one_winner( pg_factory: sessionmaker[Session], ) -> None: tenant_id, subscription_id, _ = _seed_account(pg_factory) original_payload = _cost_payload( subscription_id, idempotency_key=f"cost-{uuid.uuid4().hex}", ) with pg_factory() as db: original, _ = CommercialMeteringService(db).record_cost(tenant_id, original_payload) db.commit() original_id = original.id ready = threading.Barrier(2) def reverse_once(suffix: str) -> str: payload = original_payload.model_copy( update={ "event_type": "reversal", "idempotency_key": f"reversal-{suffix}-{uuid.uuid4().hex}", "reversal_of_cost_event_id": original_id, "occurred_at": datetime.now(UTC), } ) with pg_factory() as db: ready.wait(timeout=5) try: CommercialMeteringService(db).record_cost(tenant_id, payload) db.commit() return "created" except CommercialConflictError: db.rollback() return "blocked" with ThreadPoolExecutor(max_workers=2) as pool: outcomes = [ future.result(timeout=10) for future in (pool.submit(reverse_once, "a"), pool.submit(reverse_once, "b")) ] assert sorted(outcomes) == ["blocked", "created"] with pg_factory() as db: assert ( db.scalar( select(func.count()) .select_from(CommercialCostEvent) .where( CommercialCostEvent.tenant_id == tenant_id, CommercialCostEvent.reversal_of_cost_event_id == original_id, ) ) == 1 ) def test_concurrent_runtime_reservations_cannot_oversell_hard_quota( pg_factory: sessionmaker[Session], ) -> None: runtime_meter = { "enabled": True, "tool_type": "llm", "tool_name": "chat.completions", "quantity_basis": "call", } tenant_id, subscription_id, entitlement_id = _seed_account( pg_factory, hard_limit=Decimal("1"), runtime_meter=runtime_meter, ) ready = threading.Barrier(2) def reserve_once(suffix: str) -> str: with pg_factory() as db: ready.wait(timeout=5) try: CommercialRuntimeReservationService(db).reserve( tenant_id=tenant_id, entitlement_id=entitlement_id, run_id=f"run-{suffix}-{uuid.uuid4().hex}", tool_call_id=str(uuid.uuid4()), tool_type="llm", tool_name="chat.completions", quantity_basis="call", reserved_quantity=Decimal("1"), meter_config=runtime_meter, ) db.commit() return "reserved" except CommercialConflictError: db.rollback() return "blocked" with ThreadPoolExecutor(max_workers=2) as pool: outcomes = [ future.result(timeout=10) for future in (pool.submit(reserve_once, "a"), pool.submit(reserve_once, "b")) ] assert sorted(outcomes) == ["blocked", "reserved"] with pg_factory() as db: held = db.scalar( select(func.sum(CommercialRuntimeReservation.reserved_quantity)).where( CommercialRuntimeReservation.tenant_id == tenant_id, CommercialRuntimeReservation.subscription_id == subscription_id, CommercialRuntimeReservation.entitlement_id == entitlement_id, CommercialRuntimeReservation.status == "reserved", ) ) assert Decimal(held or 0) == Decimal("1.000000") assert ( db.scalar( select(func.count()) .select_from(UsageMeterEvent) .where(UsageMeterEvent.tenant_id == tenant_id) ) == 0 ) def test_concurrent_subscription_rollover_creates_one_period( pg_factory: sessionmaker[Session], ) -> None: tenant_id = f"tenant-rollover-{uuid.uuid4().hex}" period_start = datetime(2026, 1, 1, tzinfo=UTC) period_end = datetime(2026, 2, 1, tzinfo=UTC) with pg_factory() as db: admin = CommercialAdminService(db) plan = admin.create_plan( tenant_id, CommercialPlanCreate( plan_code="rollover", name="并发续期版", pricing_model="subscription", billing_interval="monthly", currency="CNY", base_fee=Decimal("100"), effective_from=datetime(2025, 12, 1, tzinfo=UTC), ), actor_id="postgres-test", ) admin.activate_plan(tenant_id, plan.id, expected_version=plan.version) subscription = admin.create_subscription( tenant_id, CommercialSubscriptionCreate( subscription_key=f"rollover-{uuid.uuid4().hex}", plan_id=plan.id, starts_at=period_start, current_period_start=period_start, current_period_end=period_end, seats=1, auto_renew=True, ), actor_id="postgres-test", ) db.commit() subscription_id = subscription.id ready = threading.Barrier(2) def rollover_once() -> str: with pg_factory() as db: ready.wait(timeout=5) result = CommercialSubscriptionRolloverService(db).rollover_due( tenant_id, subscription_id, as_of=period_end, ) db.commit() return result.status with ThreadPoolExecutor(max_workers=2) as pool: outcomes = [ future.result(timeout=10) for future in (pool.submit(rollover_once), pool.submit(rollover_once)) ] assert sorted(outcomes) == ["not_due", "rolled_over"] with pg_factory() as db: count = db.scalar( select(func.count()) .select_from(CommercialBillingPeriod) .where( CommercialBillingPeriod.tenant_id == tenant_id, CommercialBillingPeriod.subscription_id == subscription_id, ) ) assert count == 2 def _seed_account( factory: sessionmaker[Session], *, hard_limit: Decimal = Decimal("6"), runtime_meter: dict[str, object] | None = None, ) -> tuple[str, str, str]: tenant_id = f"tenant-commercial-{uuid.uuid4().hex}" now = datetime.now(UTC) with factory() as db: service = CommercialAdminService(db) plan = service.create_plan( tenant_id, CommercialPlanCreate( plan_code="concurrency", name="并发验证版", pricing_model="usage", billing_interval="monthly", currency="CNY", base_fee=Decimal("0"), overage_enabled=False, effective_from=now - timedelta(days=30), ), actor_id="postgres-test", ) service.activate_plan(tenant_id, plan.id, expected_version=plan.version) subscription = service.create_subscription( tenant_id, CommercialSubscriptionCreate( subscription_key=f"subscription-{uuid.uuid4().hex}", plan_id=plan.id, starts_at=now - timedelta(days=5), current_period_start=now - timedelta(days=1), current_period_end=now + timedelta(days=29), seats=1, ), actor_id="postgres-test", ) entitlement = service.upsert_entitlement( tenant_id, CommercialEntitlementUpsert( subscription_id=subscription.id, entitlement_key="concurrency_usage", metric_key="concurrency_usage", entitlement_type="metered", unit="run", included_quantity=hard_limit, hard_limit_quantity=hard_limit, reset_interval="monthly", overage_policy="block", effective_from=now - timedelta(days=5), config_json={"runtime_meter": runtime_meter} if runtime_meter else {}, ), ) db.commit() return tenant_id, subscription.id, entitlement.id def _cost_payload( subscription_id: str, *, idempotency_key: str, ) -> CommercialCostEventCreate: return CommercialCostEventCreate( subscription_id=subscription_id, cost_category="ai_inference", quantity=Decimal("10"), unit="1k_tokens", unit_cost=Decimal("1"), original_currency="CNY", reporting_currency="CNY", fx_rate=Decimal("1"), allocation_key="postgres-concurrency", occurred_at=datetime.now(UTC), source_system="postgres-provider", idempotency_key=idempotency_key, )