112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
"""
|
|
Database and Chat Injection tests for AI Chat Summaries.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
if BASE_DIR not in sys.path:
|
|
sys.path.insert(0, BASE_DIR)
|
|
|
|
from src.services.summary_db import SummaryDatabase, StoredSummary
|
|
from src.services.summary_injector import ChatSummaryInjector
|
|
|
|
|
|
class TestSummaryDatabase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_file = tempfile.NamedTemporaryFile(delete=False)
|
|
self.temp_file.close()
|
|
self.db = SummaryDatabase(self.temp_file.name)
|
|
|
|
def tearDown(self):
|
|
try:
|
|
os.unlink(self.temp_file.name)
|
|
except Exception:
|
|
pass
|
|
|
|
def test_insert_and_get_summaries(self):
|
|
summary_id = self.db.insert_summary(
|
|
dialog_id=-1001234567,
|
|
anchor_msg_id=42,
|
|
count=100,
|
|
summary_text="Discussion about system architecture and DB migrations.",
|
|
provider_name="ChatGPT (OAuth)",
|
|
model_name="terra",
|
|
offset=0,
|
|
topic_id=0,
|
|
)
|
|
self.assertTrue(summary_id > 0)
|
|
|
|
results = self.db.get_summaries_for_dialog(-1001234567, 0)
|
|
self.assertEqual(len(results), 1)
|
|
item = results[0]
|
|
self.assertEqual(item.anchor_msg_id, 42)
|
|
self.assertEqual(item.provider_name, "ChatGPT (OAuth)")
|
|
self.assertEqual(item.model_name, "terra")
|
|
self.assertIn("architecture", item.summary_text)
|
|
|
|
by_anchor = self.db.get_summary_by_anchor(-1001234567, 42)
|
|
self.assertIsNotNone(by_anchor)
|
|
self.assertEqual(by_anchor.id, summary_id)
|
|
|
|
def test_delete_summary(self):
|
|
summary_id = self.db.insert_summary(
|
|
dialog_id=555,
|
|
anchor_msg_id=10,
|
|
count=20,
|
|
summary_text="Short summary",
|
|
provider_name="Custom",
|
|
model_name="deepseek-chat",
|
|
)
|
|
self.assertTrue(self.db.delete_summary(summary_id))
|
|
self.assertIsNone(self.db.get_summary_by_anchor(555, 10))
|
|
|
|
|
|
class TestSummaryInjector(unittest.TestCase):
|
|
def test_summary_helpers_are_stable_and_plain_text(self):
|
|
from src.services.summary_injector import build_summary_message_text, summary_virtual_id
|
|
|
|
stored = StoredSummary(
|
|
id=1001,
|
|
dialog_id=-100123,
|
|
topic_id=7,
|
|
anchor_msg_id=42,
|
|
count=50,
|
|
offset=0,
|
|
summary_text="Project milestones achieved.",
|
|
provider_name="Anthropic",
|
|
model_name="claude-3-7-sonnet",
|
|
created_at=1700000000.0,
|
|
meta_json="{}",
|
|
)
|
|
self.assertEqual(summary_virtual_id(stored.id, stored.anchor_msg_id), -42001)
|
|
text = build_summary_message_text(stored)
|
|
self.assertIn("AI CHAT SUMMARY", text)
|
|
self.assertIn("Anthropic · claude-3-7-sonnet", text)
|
|
self.assertIn("Project milestones achieved.", text)
|
|
self.assertNotIn("**", text)
|
|
|
|
def test_build_client_side_message_is_safe_without_android_runtime(self):
|
|
injector = ChatSummaryInjector()
|
|
stored = StoredSummary(
|
|
id=1,
|
|
dialog_id=98765,
|
|
topic_id=0,
|
|
anchor_msg_id=100,
|
|
count=50,
|
|
offset=0,
|
|
summary_text="Project milestones achieved.",
|
|
provider_name="Anthropic",
|
|
model_name="claude-3-7-sonnet",
|
|
created_at=1700000000.0,
|
|
meta_json="{}",
|
|
)
|
|
self.assertIsNone(injector.build_client_side_summary_message(stored))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|