63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
import pytest
|
|
from ea_chatbot.history.manager import HistoryManager
|
|
from ea_chatbot.history.models import User, Conversation, Message, Plot
|
|
from ea_chatbot.config import Settings
|
|
from sqlalchemy import delete
|
|
|
|
@pytest.fixture
|
|
def history_manager():
|
|
settings = Settings()
|
|
manager = HistoryManager(settings.history_db_url)
|
|
with manager.get_session() as session:
|
|
session.execute(delete(Plot))
|
|
session.execute(delete(Message))
|
|
session.execute(delete(Conversation))
|
|
session.execute(delete(User))
|
|
return manager
|
|
|
|
def test_full_history_workflow(history_manager):
|
|
# 1. Create and Authenticate User
|
|
email = "e2e@example.com"
|
|
password = "password123"
|
|
history_manager.create_user(email, password, "E2E User")
|
|
|
|
user = history_manager.authenticate_user(email, password)
|
|
assert user is not None
|
|
assert user.display_name == "E2E User"
|
|
|
|
# 2. Create Conversation
|
|
conv = history_manager.create_conversation(user.id, "nj", "Test Analytics")
|
|
assert conv.id is not None
|
|
|
|
# 3. Add User Message
|
|
history_manager.add_message(conv.id, "user", "How many voters in NJ?")
|
|
|
|
# 4. Add Assistant Message with Plot
|
|
plot_data = b"fake_png_data"
|
|
history_manager.add_message(
|
|
conv.id,
|
|
"assistant",
|
|
"There are X voters.",
|
|
plots=[plot_data]
|
|
)
|
|
|
|
# 5. Retrieve and Verify History
|
|
messages = history_manager.get_messages(conv.id)
|
|
assert len(messages) == 2
|
|
assert messages[0].role == "user"
|
|
assert messages[1].role == "assistant"
|
|
assert len(messages[1].plots) == 1
|
|
assert messages[1].plots[0].image_data == plot_data
|
|
|
|
# 6. Verify Conversation listing
|
|
convs = history_manager.get_conversations(user.id, "nj")
|
|
assert len(convs) == 1
|
|
assert convs[0].name == "Test Analytics"
|
|
|
|
# 7. Update summary
|
|
history_manager.update_conversation_summary(conv.id, "Voter count analysis")
|
|
|
|
# 8. Reload and verify summary
|
|
updated_convs = history_manager.get_conversations(user.id, "nj")
|
|
assert updated_convs[0].summary == "Voter count analysis"
|