b23fbce8d0
Switched from username to user_id as the primary identifier in JWT tokens to better support external authentication providers. Added get_user_by_id to HistoryManager and updated API dependencies and tests to reflect these changes.
64 lines
2.7 KiB
Python
64 lines
2.7 KiB
Python
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from unittest.mock import MagicMock, patch, AsyncMock
|
|
from ea_chatbot.api.main import app
|
|
from ea_chatbot.api.dependencies import get_current_user
|
|
from ea_chatbot.history.models import User, Conversation, Message, Plot
|
|
from ea_chatbot.api.utils import create_access_token
|
|
from datetime import datetime, timezone
|
|
import json
|
|
|
|
client = TestClient(app)
|
|
|
|
@pytest.fixture
|
|
def mock_user():
|
|
return User(id="user-123", username="test@example.com", display_name="Test User")
|
|
|
|
@pytest.fixture
|
|
def auth_header(mock_user):
|
|
app.dependency_overrides[get_current_user] = lambda: mock_user
|
|
token = create_access_token(data={"sub": mock_user.id})
|
|
yield {"Authorization": f"Bearer {token}"}
|
|
app.dependency_overrides.clear()
|
|
|
|
def test_persistence_integration_success(auth_header, mock_user):
|
|
"""Test that messages and plots are persisted correctly during streaming."""
|
|
mock_events = [
|
|
{"event": "on_chat_model_stream", "name": "summarizer", "data": {"chunk": "Final answer"}},
|
|
{"event": "on_chain_end", "name": "summarizer", "data": {"output": {"messages": [{"content": "Final answer"}]}}},
|
|
{"event": "on_chain_end", "name": "summarize_conversation", "data": {"output": {"summary": "New summary"}}}
|
|
]
|
|
|
|
async def mock_astream_events(*args, **kwargs):
|
|
for event in mock_events:
|
|
yield event
|
|
|
|
with patch("ea_chatbot.api.routers.agent.app.astream_events", side_effect=mock_astream_events), \
|
|
patch("ea_chatbot.api.routers.agent.get_checkpointer") as mock_cp, \
|
|
patch("ea_chatbot.api.routers.agent.history_manager") as mock_hm:
|
|
|
|
mock_cp.return_value.__aenter__.return_value = AsyncMock()
|
|
|
|
# Mock session and DB objects
|
|
mock_session = MagicMock()
|
|
mock_hm.get_session.return_value.__enter__.return_value = mock_session
|
|
mock_conv = Conversation(id="t1", user_id=mock_user.id)
|
|
mock_session.get.return_value = mock_conv
|
|
|
|
# Act
|
|
with client.stream("POST", "/chat/stream",
|
|
json={"message": "persistence test", "thread_id": "t1"},
|
|
headers=auth_header) as response:
|
|
assert response.status_code == 200
|
|
list(response.iter_lines()) # Consume stream
|
|
|
|
# Assertions
|
|
# 1. User message should be saved immediately
|
|
mock_hm.add_message.assert_any_call("t1", "user", "persistence test")
|
|
|
|
# 2. Assistant message should be saved at the end
|
|
mock_hm.add_message.assert_any_call("t1", "assistant", "Final answer", plots=[])
|
|
|
|
# 3. Summary should be updated
|
|
mock_hm.update_conversation_summary.assert_called_once_with("t1", "New summary")
|