company-ai-platform/tests/test_subscription_dispatch.py

from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, time, timedelta
from pathlib import Path
from threading import Barrier
from uuid import uuid4

import pytest
from fastapi import HTTPException
from sqlalchemy import create_engine, func, select
from sqlalchemy.orm import Session, sessionmaker

from app.core.database import Base
from app.modules.feishu_users.constants import FeishuUserRole
from app.modules.feishu_users.models import FeishuUser
from app.modules.feishu_users.principal import FeishuPrincipal
from app.modules.subscriptions.constants import (
    DAILY_DELIVERY_LIMIT_REACHED,
    MAX_ACTIVE_SUBSCRIPTIONS,
    MAX_DAILY_DELIVERIES,
    PushDeliveryStatus,
    PushSubscriptionStatus,
    SubscriptionScheduleType,
    SubscriptionTargetType,
)
from app.modules.subscriptions.models import PushDelivery, PushSubscription
from app.modules.subscriptions.services import (
    SubscriptionManagementService,
    SubscriptionScanner,
)


def _user(
    db: Session,
    *,
    suffix: str,
    role: str = FeishuUserRole.USER,
    quiet_start: time | None = None,
    quiet_end: time | None = None,
) -> FeishuUser:
    record = FeishuUser(
        code=f"FSU-{suffix}",
        tenant_key=f"tenant-{suffix}",
        open_id=f"open-{suffix}",
        role=role,
        timezone="Asia/Shanghai",
        quiet_hours_start=quiet_start,
        quiet_hours_end=quiet_end,
    )
    db.add(record)
    db.commit()
    db.refresh(record)
    return record


def _subscription(
    db: Session,
    owner: FeishuUser,
    *,
    code: str,
    next_run_at: datetime,
    target_type: str = SubscriptionTargetType.USER,
    target_id: str | None = None,
) -> PushSubscription:
    record = PushSubscription(
        code=code,
        owner_id=owner.id,
        target_type=target_type,
        target_id=target_id or owner.open_id,
        prompt="给我一条简短提醒",
        schedule_type=SubscriptionScheduleType.DAILY,
        schedule_config={"hour": 9, "minute": 0},
        timezone=owner.timezone,
        next_run_at=next_run_at,
        status=PushSubscriptionStatus.ACTIVE,
        consented_at=next_run_at - timedelta(days=1),
    )
    db.add(record)
    db.commit()
    db.refresh(record)
    return record


def test_management_binds_private_and_group_targets() -> None:
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    try:
        with Session(engine) as db:
            user = _user(db, suffix="user")
            private_principal = FeishuPrincipal.from_user(
                user,
                chat_id="private-chat",
                chat_type="p2p",
            )
            private, _ = SubscriptionManagementService(db).create_private(
                private_principal,
                "每天 09:00",
                "给我一条提醒",
                now=datetime(2026, 7, 26, 0, 0),
            )

            assert private.target_type == SubscriptionTargetType.USER
            assert private.target_id == user.open_id

            with pytest.raises(HTTPException) as ordinary_group:
                SubscriptionManagementService(db).create_group(
                    FeishuPrincipal.from_user(
                        user,
                        chat_id="group-chat",
                        chat_type="group",
                    ),
                    "每天 10:00",
                    "群提醒",
                    now=datetime(2026, 7, 26, 0, 0),
                )
            assert ordinary_group.value.status_code == 403

            admin = _user(db, suffix="admin", role=FeishuUserRole.ADMIN)
            with pytest.raises(HTTPException, match="current group"):
                SubscriptionManagementService(db).create_group(
                    FeishuPrincipal.from_user(
                        admin,
                        chat_id="manually-supplied-chat",
                        chat_type="p2p",
                    ),
                    "每天 10:00",
                    "群提醒",
                    now=datetime(2026, 7, 26, 0, 0),
                )

            group, _ = SubscriptionManagementService(db).create_group(
                FeishuPrincipal.from_user(
                    admin,
                    chat_id="verified-current-group",
                    chat_type="group",
                ),
                "每天 10:00",
                "群提醒",
                now=datetime(2026, 7, 26, 0, 0),
            )
            assert group.target_type == SubscriptionTargetType.CHAT
            assert group.target_id == "verified-current-group"
    finally:
        engine.dispose()


def test_management_rejects_more_than_fifty_active_subscriptions() -> None:
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    try:
        with Session(engine) as db:
            user = _user(db, suffix="limit")
            for index in range(MAX_ACTIVE_SUBSCRIPTIONS):
                _subscription(
                    db,
                    user,
                    code=f"SUB-LIMIT-{index}",
                    next_run_at=datetime(2026, 7, 27, 1, 0),
                )
            with pytest.raises(HTTPException, match="at most 50"):
                SubscriptionManagementService(db).create_private(
                    FeishuPrincipal.from_user(user),
                    "每天 09:00",
                    "第 51 条",
                    now=datetime(2026, 7, 26, 0, 0),
                )
    finally:
        engine.dispose()


def test_scanner_materializes_one_window_and_advances_the_plan() -> None:
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    current = datetime(2026, 7, 26, 1, 0)
    try:
        with Session(engine) as db:
            user = _user(db, suffix="scan")
            subscription = _subscription(
                db,
                user,
                code="SUB-SCAN",
                next_run_at=current,
            )

            first = SubscriptionScanner(db).scan_due(now=current)
            second = SubscriptionScanner(db).scan_due(now=current)

            assert len(first) == 1
            assert second == []
            assert first[0].scheduled_for == current
            assert first[0].status == PushDeliveryStatus.PENDING
            assert first[0].next_attempt_at == current
            db.refresh(subscription)
            assert subscription.next_run_at == datetime(2026, 7, 27, 1, 0)
            assert db.scalar(select(func.count()).select_from(PushDelivery)) == 1
    finally:
        engine.dispose()


def test_scanner_defers_quiet_hours_and_skips_daily_limit() -> None:
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    current = datetime(2026, 7, 26, 15, 30)
    try:
        with Session(engine) as db:
            quiet_user = _user(
                db,
                suffix="quiet",
                quiet_start=time(22, 0),
                quiet_end=time(7, 0),
            )
            _subscription(
                db,
                quiet_user,
                code="SUB-QUIET",
                next_run_at=current,
            )
            quiet_delivery = SubscriptionScanner(db).scan_due(now=current)[0]
            assert quiet_delivery.status == PushDeliveryStatus.PENDING
            assert quiet_delivery.next_attempt_at == datetime(2026, 7, 26, 23, 0)

            limited_user = _user(db, suffix="daily-limit")
            history = _subscription(
                db,
                limited_user,
                code="SUB-HISTORY",
                next_run_at=current + timedelta(days=1),
            )
            for index in range(MAX_DAILY_DELIVERIES):
                db.add(
                    PushDelivery(
                        code=f"DEL-HISTORY-{index}",
                        subscription_id=history.id,
                        scheduled_for=current - timedelta(minutes=index),
                        idempotency_key=f"{index:064x}",
                        message_uuid=str(uuid4()),
                        status=PushDeliveryStatus.SENT,
                        attempt_count=1,
                        sent_at=current,
                        created_at=current,
                        updated_at=current,
                    )
                )
            _subscription(
                db,
                limited_user,
                code="SUB-OVER-LIMIT",
                next_run_at=current,
            )
            db.commit()

            limited_delivery = SubscriptionScanner(db).scan_due(now=current)[0]
            assert limited_delivery.status == PushDeliveryStatus.SKIPPED
            assert limited_delivery.last_error == DAILY_DELIVERY_LIMIT_REACHED
    finally:
        engine.dispose()


def test_two_sqlite_scanners_claim_only_one_delivery(tmp_path: Path) -> None:
    database_path = tmp_path / "subscription-scanner.db"
    engine = create_engine(
        f"sqlite:///{database_path}",
        connect_args={"check_same_thread": False, "timeout": 10},
    )
    with engine.begin() as connection:
        connection.exec_driver_sql("PRAGMA journal_mode=WAL")
    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)
    current = datetime(2026, 7, 26, 1, 0)
    try:
        with factory() as db:
            user = _user(db, suffix="concurrent")
            _subscription(
                db,
                user,
                code="SUB-CONCURRENT",
                next_run_at=current,
            )

        barrier = Barrier(2)

        def scan(worker_id: str) -> list[str]:
            with factory() as db:
                barrier.wait()
                return [
                    delivery.code
                    for delivery in SubscriptionScanner(db).scan_due(
                        now=current,
                        worker_id=worker_id,
                    )
                ]

        with ThreadPoolExecutor(max_workers=2) as executor:
            results = list(executor.map(scan, ("scanner-one", "scanner-two")))

        assert sum(len(items) for items in results) == 1
        with factory() as db:
            assert db.scalar(select(func.count()).select_from(PushDelivery)) == 1
    finally:
        engine.dispose()