146 lines
5.1 KiB
Python
146 lines
5.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)
|
|
# Clean up tables before tests (order matters because of foreign keys)
|
|
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_history_manager_initialization(history_manager):
|
|
assert history_manager.engine is not None
|
|
assert history_manager.SessionLocal is not None
|
|
|
|
def test_history_manager_session_context(history_manager):
|
|
with history_manager.get_session() as session:
|
|
assert session is not None
|
|
|
|
def test_get_user_not_found(history_manager):
|
|
user = history_manager.get_user("nonexistent@example.com")
|
|
assert user is None
|
|
|
|
def test_authenticate_user_success(history_manager):
|
|
email = "test@example.com"
|
|
password = "secretpassword"
|
|
history_manager.create_user(email=email, password=password)
|
|
|
|
user = history_manager.authenticate_user(email, password)
|
|
assert user is not None
|
|
assert user.username == email
|
|
|
|
def test_authenticate_user_failure(history_manager):
|
|
email = "test@example.com"
|
|
history_manager.create_user(email=email, password="correctpassword")
|
|
|
|
user = history_manager.authenticate_user(email, "wrongpassword")
|
|
assert user is None
|
|
|
|
def test_sync_user_from_oidc_new_user(history_manager):
|
|
user = history_manager.sync_user_from_oidc(
|
|
email="new@example.com",
|
|
display_name="New User"
|
|
)
|
|
assert user is not None
|
|
assert user.username == "new@example.com"
|
|
assert user.display_name == "New User"
|
|
|
|
def test_sync_user_from_oidc_existing_user(history_manager):
|
|
# First sync
|
|
history_manager.sync_user_from_oidc(
|
|
email="existing@example.com",
|
|
display_name="First Name"
|
|
)
|
|
# Second sync should update or return same user
|
|
user = history_manager.sync_user_from_oidc(
|
|
email="existing@example.com",
|
|
display_name="Updated Name"
|
|
)
|
|
assert user.display_name == "Updated Name"
|
|
|
|
# --- Conversation Management Tests ---
|
|
|
|
@pytest.fixture
|
|
def user(history_manager):
|
|
return history_manager.create_user(email="conv_user@example.com")
|
|
|
|
def test_create_conversation(history_manager, user):
|
|
conv = history_manager.create_conversation(
|
|
user_id=user.id,
|
|
data_state="new_jersey",
|
|
name="Test Chat",
|
|
summary="A test conversation summary"
|
|
)
|
|
assert conv is not None
|
|
assert conv.name == "Test Chat"
|
|
assert conv.summary == "A test conversation summary"
|
|
assert conv.user_id == user.id
|
|
|
|
def test_get_conversations(history_manager, user):
|
|
history_manager.create_conversation(user_id=user.id, data_state="nj", name="C1")
|
|
history_manager.create_conversation(user_id=user.id, data_state="nj", name="C2")
|
|
history_manager.create_conversation(user_id=user.id, data_state="ny", name="C3")
|
|
|
|
nj_convs = history_manager.get_conversations(user_id=user.id, data_state="nj")
|
|
assert len(nj_convs) == 2
|
|
|
|
ny_convs = history_manager.get_conversations(user_id=user.id, data_state="ny")
|
|
assert len(ny_convs) == 1
|
|
|
|
def test_rename_conversation(history_manager, user):
|
|
conv = history_manager.create_conversation(user.id, "nj", "Old Name")
|
|
updated = history_manager.rename_conversation(conv.id, "New Name")
|
|
assert updated.name == "New Name"
|
|
|
|
def test_delete_conversation(history_manager, user):
|
|
conv = history_manager.create_conversation(user.id, "nj", "To Delete")
|
|
history_manager.delete_conversation(conv.id)
|
|
|
|
convs = history_manager.get_conversations(user.id, "nj")
|
|
assert len(convs) == 0
|
|
|
|
# --- Message Management Tests ---
|
|
|
|
@pytest.fixture
|
|
def conversation(history_manager, user):
|
|
return history_manager.create_conversation(user.id, "nj", "Msg Test Conv")
|
|
|
|
def test_add_message(history_manager, conversation):
|
|
msg = history_manager.add_message(
|
|
conversation_id=conversation.id,
|
|
role="user",
|
|
content="Hello world"
|
|
)
|
|
assert msg is not None
|
|
assert msg.content == "Hello world"
|
|
assert msg.role == "user"
|
|
assert msg.conversation_id == conversation.id
|
|
|
|
def test_add_message_with_plots(history_manager, conversation):
|
|
plots_data = [b"fake_plot_1", b"fake_plot_2"]
|
|
msg = history_manager.add_message(
|
|
conversation_id=conversation.id,
|
|
role="assistant",
|
|
content="Here are plots",
|
|
plots=plots_data
|
|
)
|
|
assert len(msg.plots) == 2
|
|
assert msg.plots[0].image_data == b"fake_plot_1"
|
|
|
|
def test_get_messages(history_manager, conversation):
|
|
history_manager.add_message(conversation.id, "user", "Q1")
|
|
history_manager.add_message(conversation.id, "assistant", "A1")
|
|
|
|
messages = history_manager.get_messages(conversation.id)
|
|
assert len(messages) == 2
|
|
assert messages[0].content == "Q1"
|
|
assert messages[1].content == "A1"
|