91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from datetime import UTC, datetime
|
||
|
|
from decimal import Decimal
|
||
|
|
|
||
|
|
from sqlalchemy import create_engine, func, select
|
||
|
|
from sqlalchemy.orm import Session, sessionmaker
|
||
|
|
from sqlalchemy.pool import StaticPool
|
||
|
|
|
||
|
|
import app.models # noqa: F401 - 注册完整 metadata
|
||
|
|
from app.db.base_class import Base
|
||
|
|
from app.models.commercial_billing import CommercialBillingPeriod
|
||
|
|
from app.schemas.commercial import CommercialPlanCreate, CommercialSubscriptionCreate
|
||
|
|
from app.services.commercial_admin import CommercialAdminService
|
||
|
|
from app.services.commercial_rollover_scheduler import CommercialRolloverScheduler
|
||
|
|
|
||
|
|
|
||
|
|
def test_scheduler_rolls_due_subscription_once_and_respects_leader_lease() -> None:
|
||
|
|
engine = create_engine(
|
||
|
|
"sqlite+pysqlite:///:memory:",
|
||
|
|
connect_args={"check_same_thread": False},
|
||
|
|
poolclass=StaticPool,
|
||
|
|
)
|
||
|
|
Base.metadata.create_all(engine)
|
||
|
|
factory = sessionmaker(bind=engine, expire_on_commit=False)
|
||
|
|
due = datetime(2026, 2, 1, tzinfo=UTC)
|
||
|
|
with factory() as db:
|
||
|
|
subscription_id = _seed_due_subscription(db)
|
||
|
|
db.commit()
|
||
|
|
|
||
|
|
scheduler = CommercialRolloverScheduler(session_factory=factory)
|
||
|
|
first = scheduler._run_once(as_of=due)
|
||
|
|
replay = scheduler._run_once(as_of=due)
|
||
|
|
|
||
|
|
assert first["scanned"] == 1
|
||
|
|
assert first["rolled_over"] == 1
|
||
|
|
assert first["periods_created"] == 1
|
||
|
|
assert replay["scanned"] == 0
|
||
|
|
with factory() as db:
|
||
|
|
assert (
|
||
|
|
db.scalar(
|
||
|
|
select(func.count())
|
||
|
|
.select_from(CommercialBillingPeriod)
|
||
|
|
.where(CommercialBillingPeriod.subscription_id == subscription_id)
|
||
|
|
)
|
||
|
|
== 2
|
||
|
|
)
|
||
|
|
|
||
|
|
scheduler._try_acquire_lease = lambda _db: False # type: ignore[method-assign]
|
||
|
|
skipped = scheduler._run_once(as_of=due)
|
||
|
|
assert skipped["leader_skipped"] == 1
|
||
|
|
engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
def _seed_due_subscription(db: Session) -> str:
|
||
|
|
start = datetime(2026, 1, 1, tzinfo=UTC)
|
||
|
|
due = datetime(2026, 2, 1, tzinfo=UTC)
|
||
|
|
admin = CommercialAdminService(db)
|
||
|
|
plan = admin.create_plan(
|
||
|
|
"tenant-scheduler",
|
||
|
|
CommercialPlanCreate(
|
||
|
|
plan_code="scheduler",
|
||
|
|
name="调度器测试套餐",
|
||
|
|
pricing_model="subscription",
|
||
|
|
billing_interval="monthly",
|
||
|
|
currency="CNY",
|
||
|
|
base_fee=Decimal("100"),
|
||
|
|
effective_from=start,
|
||
|
|
),
|
||
|
|
actor_id="scheduler-test",
|
||
|
|
)
|
||
|
|
admin.activate_plan(
|
||
|
|
"tenant-scheduler",
|
||
|
|
plan.id,
|
||
|
|
expected_version=plan.version,
|
||
|
|
)
|
||
|
|
subscription = admin.create_subscription(
|
||
|
|
"tenant-scheduler",
|
||
|
|
CommercialSubscriptionCreate(
|
||
|
|
subscription_key="scheduler-subscription",
|
||
|
|
plan_id=plan.id,
|
||
|
|
starts_at=start,
|
||
|
|
current_period_start=start,
|
||
|
|
current_period_end=due,
|
||
|
|
seats=1,
|
||
|
|
auto_renew=True,
|
||
|
|
),
|
||
|
|
actor_id="scheduler-test",
|
||
|
|
)
|
||
|
|
return subscription.id
|