feat: add entitlement plans and topic sessions

ai_knowledge_base_v2/apps/backend/tests/test_entitlements_topics.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+import pytest
+from fastapi import HTTPException
+from sqlalchemy import create_engine
+from sqlalchemy.orm import Session
+from sqlalchemy.pool import StaticPool
+
+from app.models import Base
+from app.models.chat import ChatSession, TopicSession
+from app.models.entitlement import EntitlementPlan
+from app.models.user import User
+from app.services.chat_service import ChatService
+from app.services.entitlement_service import EntitlementService
+from app.services.topic_session_service import TopicSessionService
+
+
+def _db() -> Session:
+    engine = create_engine("sqlite:///:memory:", connect_args={"check_same_thread": False}, poolclass=StaticPool)
+    Base.metadata.create_all(engine)
+    return Session(engine)
+
+
+def _now() -> datetime:
+    return datetime.now(UTC).replace(tzinfo=None)
+
+
+def _seed_user_session(db: Session) -> tuple[User, ChatSession]:
+    user = User(id=1, phone="13800000001", name="测试用户", daily_chat_limit=100, daily_chat_used=0)
+    session = ChatSession(id=1, user_id=1, title="新聊天", message_count=0, last_message_at=_now(), is_deleted=0)
+    db.add_all([user, session])
+    db.commit()
+    return user, session
+
+
+def test_default_entitlement_uses_basic_plan_when_user_has_no_assignment():
+    with _db() as db:
+        user, _session = _seed_user_session(db)
+        db.add(
+            EntitlementPlan(
+                id=10,
+                name="大本营基础版",
+                plan_type="basic",
+                monthly_topic_limit=30,
+                status=1,
+                sort_order=10,
+            )
+        )
+        db.commit()
+
+        view = EntitlementService.active_entitlement(db, user, monthly_topic_used=3)
+
+        assert view.plan_id == 10
+        assert view.name == "大本营基础版"
+        assert view.source == "default"
+        assert view.monthly_topic_remaining == 27
+
+
+def test_assign_user_plan_replaces_previous_active_plan():
+    with _db() as db:
+        user, _session = _seed_user_session(db)
+        db.add_all(
+            [
+                EntitlementPlan(id=10, name="基础版", plan_type="basic", monthly_topic_limit=30, status=1, sort_order=10),
+                EntitlementPlan(id=20, name="深度陪伴版", plan_type="deep", monthly_topic_limit=90, status=1, sort_order=20),
+            ]
+        )
+        db.commit()
+
+        first = EntitlementService.assign_user_plan(db, user=user, plan_id=10, operated_by=99)
+        second = EntitlementService.assign_user_plan(db, user=user, plan_id=20, operated_by=99)
+        db.commit()
+
+        db.refresh(first)
+        db.refresh(second)
+        view = EntitlementService.active_entitlement(db, user, monthly_topic_used=4)
+
+        assert first.status == "replaced"
+        assert second.status == "active"
+        assert view.plan_id == 20
+        assert view.source == "assigned"
+        assert view.monthly_topic_remaining == 86
+
+
+def test_monthly_topic_quota_blocks_new_topic_but_allows_existing_topic():
+    with _db() as db:
+        user, session = _seed_user_session(db)
+        db.add(EntitlementPlan(id=10, name="限额版", plan_type="basic", monthly_topic_limit=1, status=1, sort_order=10))
+        db.add(
+            TopicSession(
+                id=100,
+                user_id=user.id,
+                chat_session_id=99,
+                title="旧主题",
+                core_question="旧主题",
+                status="active",
+                quota_deducted=1,
+                started_at=_now(),
+                created_at=_now(),
+                updated_at=_now(),
+            )
+        )
+        db.commit()
+        entitlement = EntitlementService.active_entitlement(db, user, monthly_topic_used=TopicSessionService.monthly_used_count(db, user.id))
+
+        with pytest.raises(HTTPException) as exc:
+            ChatService._ensure_topic_quota(db, user, session, entitlement)
+
+        assert exc.value.status_code == 403
+
+        TopicSessionService.get_or_create_active(
+            db,
+            user=user,
+            session=session,
+            question="当前主题",
+            deduct_quota=True,
+        )
+        db.flush()
+
+        ChatService._ensure_topic_quota(db, user, session, entitlement)