128 lines
5.5 KiB
Python
128 lines
5.5 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from decimal import Decimal
|
|
from pathlib import Path
|
|
|
|
from agent_integration.client import OpenAgentEvent
|
|
from agent_integration.service import OpenAgentService, SessionNotFoundError
|
|
from agent_integration.sessions import SQLiteSessionStore
|
|
|
|
|
|
class FakeAgentClient:
|
|
def __init__(self):
|
|
self.create_calls = []
|
|
self.send_calls = []
|
|
self.stream_calls = []
|
|
self.run_calls = []
|
|
|
|
def create_session(self, **kwargs):
|
|
self.create_calls.append(kwargs)
|
|
return {"session_id": "open_sess_123", "status": "active"}
|
|
|
|
def send_message(self, session_id, message, **kwargs):
|
|
self.send_calls.append((session_id, message, kwargs))
|
|
return {"session_id": session_id, "run_id": "run_123", "status": "running"}
|
|
|
|
def stream_message(self, session_id, message, **kwargs):
|
|
self.stream_calls.append((session_id, message, kwargs))
|
|
yield OpenAgentEvent(event="message.delta", data={"content": "你"})
|
|
|
|
def get_run(self, session_id, run_id):
|
|
self.run_calls.append(("get", session_id, run_id))
|
|
return {"session_id": session_id, "run_id": run_id, "status": "completed"}
|
|
|
|
def cancel_run(self, session_id, run_id):
|
|
self.run_calls.append(("cancel", session_id, run_id))
|
|
return {"session_id": session_id, "run_id": run_id, "status": "cancelling"}
|
|
|
|
|
|
class AgentServiceTests(unittest.TestCase):
|
|
def test_session_is_created_once_and_persisted(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "sessions.sqlite3"
|
|
client = FakeAgentClient()
|
|
with SQLiteSessionStore(db_path) as store:
|
|
service = OpenAgentService(client, store) # type: ignore[arg-type]
|
|
first = service.ensure_session(
|
|
"conversation-001",
|
|
external_subject_id="customer-001",
|
|
metadata={"source": "unittest"},
|
|
)
|
|
second = service.ensure_session("conversation-001")
|
|
|
|
with SQLiteSessionStore(db_path) as reopened_store:
|
|
persisted = reopened_store.get("conversation-001")
|
|
|
|
self.assertEqual(first, second)
|
|
self.assertEqual(first, persisted)
|
|
self.assertEqual(len(client.create_calls), 1)
|
|
session_key = client.create_calls[0]["idempotency_key"]
|
|
self.assertTrue(session_key.startswith("sess_"))
|
|
self.assertNotIn("conversation-001", session_key)
|
|
|
|
def test_stable_message_id_produces_stable_idempotency_key(self):
|
|
client = FakeAgentClient()
|
|
with SQLiteSessionStore(":memory:") as store:
|
|
service = OpenAgentService(client, store) # type: ignore[arg-type]
|
|
service.send_message("conversation-001", "first", message_id="upstream-message-001")
|
|
service.send_message("conversation-001", "retry", message_id="upstream-message-001")
|
|
|
|
self.assertEqual(len(client.create_calls), 1)
|
|
first_key = client.send_calls[0][2]["idempotency_key"]
|
|
second_key = client.send_calls[1][2]["idempotency_key"]
|
|
self.assertEqual(first_key, second_key)
|
|
self.assertTrue(first_key.startswith("msg_"))
|
|
self.assertNotIn("upstream-message-001", first_key)
|
|
|
|
def test_stream_get_cancel_and_forget_use_saved_session(self):
|
|
client = FakeAgentClient()
|
|
with SQLiteSessionStore(":memory:") as store:
|
|
service = OpenAgentService(client, store) # type: ignore[arg-type]
|
|
events = list(service.stream_message("conversation-001", "你好", message_id="message-001"))
|
|
status = service.get_run("conversation-001", "run-001")
|
|
cancelled = service.cancel_run("conversation-001", "run-001")
|
|
deleted = service.forget_session("conversation-001")
|
|
with self.assertRaises(SessionNotFoundError):
|
|
service.get_session("conversation-001")
|
|
|
|
self.assertEqual(events, [OpenAgentEvent(event="message.delta", data={"content": "你"})])
|
|
self.assertEqual(status["status"], "completed")
|
|
self.assertEqual(cancelled["status"], "cancelling")
|
|
self.assertTrue(deleted)
|
|
self.assertEqual(
|
|
client.run_calls,
|
|
[
|
|
("get", "open_sess_123", "run-001"),
|
|
("cancel", "open_sess_123", "run-001"),
|
|
],
|
|
)
|
|
|
|
def test_unknown_conversation_fails_without_remote_call(self):
|
|
client = FakeAgentClient()
|
|
with SQLiteSessionStore(":memory:") as store:
|
|
service = OpenAgentService(client, store) # type: ignore[arg-type]
|
|
with self.assertRaises(SessionNotFoundError):
|
|
service.get_run("missing-conversation", "run-001")
|
|
self.assertEqual(client.run_calls, [])
|
|
|
|
def test_message_whitespace_is_preserved_and_metadata_is_strict_json(self):
|
|
client = FakeAgentClient()
|
|
with SQLiteSessionStore(":memory:") as store:
|
|
service = OpenAgentService(client, store) # type: ignore[arg-type]
|
|
service.send_message("conversation-001", " keep spacing ", message_id="message-001")
|
|
with self.assertRaises(ValueError):
|
|
service.send_message(
|
|
"conversation-002",
|
|
"hello",
|
|
metadata={"not_json": Decimal("1.5")},
|
|
)
|
|
|
|
self.assertEqual(client.send_calls[0][1], " keep spacing ")
|
|
self.assertEqual(len(client.create_calls), 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|