mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
zeekay
ea3116aac9
- Remove all unconditional pytest.skip() calls - Fix file content tests to handle both strict/non-strict modes - Fix missing device_capabilities module in hanzo-network - Fix HanzoNetProvider imports and exports - Fix local_llm.py import errors - Create test helper classes for pagination tests - Fix FastMCP paginator test constructor args - Configure all API keys from ~/.zsh/after.zsh - All tests now pass or skip conditionally based on actual requirements
231 lines
7.6 KiB
Python
231 lines
7.6 KiB
Python
"""Tests for Chat API endpoints."""
|
|
|
|
import polars as pl
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from hanzo_memory.server import app
|
|
|
|
|
|
class TestChatAPI:
|
|
"""Test Chat API endpoints."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup(self, mock_auth, mock_db_client, mock_services):
|
|
"""Set up test client and mocks."""
|
|
self.client = TestClient(app)
|
|
self.headers = {"Authorization": "Bearer test-token"}
|
|
self.mock_db = mock_db_client
|
|
self.mock_embedding_service = mock_services["embedding"]
|
|
|
|
def test_create_chat_session(self):
|
|
"""Test creating a chat session."""
|
|
response = self.client.post(
|
|
"/v1/chat/sessions/create",
|
|
json={
|
|
"userid": "user123",
|
|
"session_id": "session123",
|
|
"project_id": "proj123",
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["session_id"] == "session123"
|
|
assert data["userid"] == "user123"
|
|
assert data["project_id"] == "proj123"
|
|
assert data["created"] is True
|
|
|
|
# Verify chat table was created
|
|
self.mock_db.create_chats_table.assert_called_once_with("user123")
|
|
|
|
def test_create_chat_session_auto_ids(self):
|
|
"""Test creating a chat session with auto-generated IDs."""
|
|
response = self.client.post(
|
|
"/v1/chat/sessions/create",
|
|
json={"userid": "user123"},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "session_id" in data
|
|
assert data["userid"] == "user123"
|
|
assert data["project_id"] == "default-user123"
|
|
|
|
def test_add_chat_message(self):
|
|
"""Test adding a chat message."""
|
|
# Mock empty search results (no duplicates)
|
|
self.mock_db.search_chats.return_value = pl.DataFrame()
|
|
self.mock_db.add_chat_message.return_value = {"chat_id": "msg123"}
|
|
|
|
response = self.client.post(
|
|
"/v1/chat/messages/add",
|
|
json={
|
|
"userid": "user123",
|
|
"session_id": "session123",
|
|
"role": "user",
|
|
"content": "Hello, how are you?",
|
|
"metadata": {"timestamp": "2023-01-01T12:00:00Z"},
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "chat_id" in data
|
|
assert data["session_id"] == "session123"
|
|
assert data["duplicate"] is False
|
|
|
|
# Verify embedding was generated
|
|
self.mock_embedding_service.embed_text.assert_called_once_with(
|
|
"Hello, how are you?"
|
|
)
|
|
|
|
# Verify message was added
|
|
self.mock_db.add_chat_message.assert_called_once()
|
|
|
|
def test_add_duplicate_chat_message(self):
|
|
"""Test adding a duplicate chat message."""
|
|
# Mock search results with duplicate
|
|
mock_df = pl.DataFrame(
|
|
{
|
|
"chat_id": ["existing123"],
|
|
"content": ["Hello, how are you?"],
|
|
"role": ["user"],
|
|
"_similarity": [0.995],
|
|
}
|
|
)
|
|
self.mock_db.search_chats.return_value = mock_df
|
|
|
|
response = self.client.post(
|
|
"/v1/chat/messages/add",
|
|
json={
|
|
"userid": "user123",
|
|
"session_id": "session123",
|
|
"role": "user",
|
|
"content": "Hello, how are you?",
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["chat_id"] == "existing123"
|
|
assert data["duplicate"] is True
|
|
|
|
# Verify message was NOT added
|
|
self.mock_db.add_chat_message.assert_not_called()
|
|
|
|
def test_get_chat_messages(self):
|
|
"""Test getting chat messages for a session."""
|
|
# Mock chat history
|
|
mock_df = pl.DataFrame(
|
|
{
|
|
"chat_id": ["msg1", "msg2", "msg3"],
|
|
"role": ["user", "assistant", "user"],
|
|
"content": ["Hello", "Hi there!", "How are you?"],
|
|
"metadata": ["{}", '{"model": "gpt-4"}', "{}"],
|
|
"created_at": [
|
|
"2023-01-01T12:00:00",
|
|
"2023-01-01T12:01:00",
|
|
"2023-01-01T12:02:00",
|
|
],
|
|
}
|
|
)
|
|
self.mock_db.get_chat_history.return_value = mock_df
|
|
|
|
response = self.client.get(
|
|
"/v1/chat/sessions/session123/messages",
|
|
params={"userid": "user123", "limit": 100},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["session_id"] == "session123"
|
|
assert len(data["messages"]) == 3
|
|
assert data["total"] == 3
|
|
|
|
# Check message order
|
|
assert data["messages"][0]["content"] == "Hello"
|
|
assert data["messages"][1]["content"] == "Hi there!"
|
|
assert data["messages"][2]["content"] == "How are you?"
|
|
|
|
def test_search_chat_messages(self):
|
|
"""Test searching chat messages."""
|
|
# Mock search results
|
|
mock_df = pl.DataFrame(
|
|
{
|
|
"chat_id": ["msg1", "msg2"],
|
|
"session_id": ["session1", "session2"],
|
|
"role": ["user", "assistant"],
|
|
"content": ["Tell me about Python", "Python is a programming language"],
|
|
"_similarity": [0.95, 0.92],
|
|
"created_at": ["2023-01-01T12:00:00", "2023-01-01T13:00:00"],
|
|
}
|
|
)
|
|
self.mock_db.search_chats.return_value = mock_df
|
|
|
|
response = self.client.post(
|
|
"/v1/chat/search",
|
|
params={
|
|
"query": "Python programming",
|
|
"userid": "user123",
|
|
"limit": 10,
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["query"] == "Python programming"
|
|
assert len(data["messages"]) == 2
|
|
assert data["messages"][0]["similarity_score"] == 0.95
|
|
assert "Python" in data["messages"][0]["content"]
|
|
|
|
# Verify search was called
|
|
self.mock_db.search_chats.assert_called_once()
|
|
search_args = self.mock_db.search_chats.call_args[1]
|
|
assert search_args["user_id"] == "user123"
|
|
assert search_args["limit"] == 10
|
|
|
|
def test_search_chat_with_filters(self):
|
|
"""Test searching chat messages with filters."""
|
|
self.mock_db.search_chats.return_value = pl.DataFrame()
|
|
|
|
response = self.client.post(
|
|
"/v1/chat/search",
|
|
params={
|
|
"query": "test query",
|
|
"userid": "user123",
|
|
"project_id": "proj123",
|
|
"session_id": "session123",
|
|
"limit": 5,
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
# Verify filters were passed
|
|
search_args = self.mock_db.search_chats.call_args[1]
|
|
assert search_args["project_id"] == "proj123"
|
|
assert search_args["session_id"] == "session123"
|
|
|
|
def test_chat_error_handling(self):
|
|
"""Test error handling in chat operations."""
|
|
self.mock_db.create_chats_table.side_effect = Exception("DB Error")
|
|
|
|
response = self.client.post(
|
|
"/v1/chat/sessions/create",
|
|
json={"userid": "user123"},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 500
|
|
data = response.json()
|
|
assert "detail" in data
|
|
assert "DB Error" in data["detail"]
|