john
05e173b293
Replace the WebView shell with SwiftUI screens, add account-scoped Wiki and TTS APIs with adaptive review and photo scan support, and keep web/iOS pages usable while data loads asynchronously. Co-authored-by: Cursor <cursoragent@cursor.com>
61 lines
2.2 KiB
Python
61 lines
2.2 KiB
Python
import os
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from database import Base
|
|
from models import User, WikiLlmUsage
|
|
from services.wiki_llm_quota_service import WikiLlmQuotaExceeded, wiki_llm_quota_service
|
|
|
|
|
|
class WikiLlmQuotaServiceTest(unittest.TestCase):
|
|
def setUp(self):
|
|
engine = create_engine("sqlite:///:memory:")
|
|
Base.metadata.create_all(engine)
|
|
self.db = sessionmaker(bind=engine)()
|
|
self.user = User(
|
|
username="alice",
|
|
password_hash="test",
|
|
created_at="2026-06-06T00:00:00Z",
|
|
)
|
|
self.db.add(self.user)
|
|
self.db.commit()
|
|
|
|
def tearDown(self):
|
|
self.db.close()
|
|
get_settings = __import__("config", fromlist=["get_settings"]).get_settings
|
|
get_settings.cache_clear()
|
|
|
|
def test_consume_until_limit(self):
|
|
with patch.dict(os.environ, {"AI_WIKI_DAILY_LLM_LIMIT": "2"}, clear=False):
|
|
__import__("config", fromlist=["get_settings"]).get_settings.cache_clear()
|
|
for _ in range(2):
|
|
wiki_llm_quota_service.consume(self.db, self.user)
|
|
self.db.commit()
|
|
with self.assertRaises(WikiLlmQuotaExceeded):
|
|
wiki_llm_quota_service.consume(self.db, self.user)
|
|
|
|
def test_quota_dict_tracks_usage(self):
|
|
with patch.dict(os.environ, {"AI_WIKI_DAILY_LLM_LIMIT": "20"}, clear=False):
|
|
__import__("config", fromlist=["get_settings"]).get_settings.cache_clear()
|
|
wiki_llm_quota_service.consume(self.db, self.user)
|
|
self.db.commit()
|
|
quota = wiki_llm_quota_service.quota_dict(self.db, self.user)
|
|
self.assertEqual(quota["used_today"], 1)
|
|
self.assertEqual(quota["remaining_today"], 19)
|
|
self.assertFalse(quota["quota_exceeded"])
|
|
|
|
def test_usage_row_is_unique_per_day(self):
|
|
row = wiki_llm_quota_service.get_usage_row(self.db, self.user, "2026-06-06")
|
|
row.call_count = 3
|
|
self.db.commit()
|
|
rows = self.db.query(WikiLlmUsage).filter(WikiLlmUsage.user_id == self.user.id).all()
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(rows[0].call_count, 3)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|