import pytest
# Test prompts and data
CONTEXT_PROMPT = "Who is the CEO of Microsoft?"
FOLLOW_UP_PROMPT = "When was he born?"
RAG_DOC_TITLE = "Fictional Company History"
RAG_DOC_TEXT = "The company AlphaCorp was founded in 2021 by Jane Doe. Their primary product is a smart home device called 'Nexus'."
RAG_PROMPT = "Who founded AlphaCorp and what is their main product?"
@pytest.mark.asyncio
async def test_chat_in_session_lifecycle(http_client):
"""
Tests a full session lifecycle from creation to conversational memory.
This test is a single, sequential unit.
"""
print("\n--- Running test_chat_in_session_lifecycle ---")
# 1. Create a new session with a trailing slash
payload = {"user_id": "integration_tester_lifecycle", "provider_name": "deepseek"}
response = await http_client.post("/sessions/", json=payload)
assert response.status_code == 200
session_id = response.json()["id"]
print(f"✅ Session created successfully with ID: {session_id}")
# 2. First chat turn to establish context
chat_payload_1 = {"prompt": CONTEXT_PROMPT}
response_1 = await http_client.post(f"/sessions/{session_id}/chat", json=chat_payload_1)
assert response_1.status_code == 200
assert "Satya Nadella" in response_1.json()["answer"]
assert response_1.json()["provider_used"] == "deepseek"
print("✅ Chat Turn 1 (context) test passed.")
# 3. Second chat turn (follow-up) to test conversational memory
chat_payload_2 = {"prompt": FOLLOW_UP_PROMPT}
response_2 = await http_client.post(f"/sessions/{session_id}/chat", json=chat_payload_2)
assert response_2.status_code == 200
assert "1967" in response_2.json()["answer"]
assert response_2.json()["provider_used"] == "deepseek"
print("✅ Chat Turn 2 (follow-up) test passed.")
# 4. Cleanup (optional, but good practice if not using a test database that resets)
# The session data would typically be cleaned up by the database teardown.
@pytest.mark.asyncio
async def test_chat_with_model_switch(http_client, session_id):
"""Tests switching models within an existing session."""
print("\n--- Running test_chat_with_model_switch ---")
# Send a message to the new session with a different model
payload_gemini = {"prompt": "What is the capital of France?", "provider_name": "gemini"}
response_gemini = await http_client.post(f"/sessions/{session_id}/chat", json=payload_gemini)
assert response_gemini.status_code == 200
assert "Paris" in response_gemini.json()["answer"]
assert response_gemini.json()["provider_used"] == "gemini"
print("✅ Chat (Model Switch to Gemini) test passed.")
# Switch back to the original model
payload_deepseek = {"prompt": "What is the largest ocean?", "provider_name": "deepseek"}
response_deepseek = await http_client.post(f"/sessions/{session_id}/chat", json=payload_deepseek)
assert response_deepseek.status_code == 200
assert "Pacific Ocean" in response_deepseek.json()["answer"]
assert response_deepseek.json()["provider_used"] == "deepseek"
print("✅ Chat (Model Switch back to DeepSeek) test passed.")
@pytest.mark.asyncio
async def test_chat_with_document_retrieval(http_client):
"""
Tests injecting a document and using it for retrieval-augmented generation.
This test creates its own session and document for isolation.
"""
print("\n--- Running test_chat_with_document_retrieval ---")
# Create a new session for this RAG test
session_response = await http_client.post("/sessions/", json={"user_id": "rag_tester", "provider_name": "deepseek"})
assert session_response.status_code == 200
rag_session_id = session_response.json()["id"]
# Add a new document with specific content for retrieval
doc_data = {"title": RAG_DOC_TITLE, "text": RAG_DOC_TEXT}
add_doc_response = await http_client.post("/documents/", json=doc_data)
assert add_doc_response.status_code == 200
try:
message = add_doc_response.json().get("message", "")
rag_document_id = int(message.split(" with ID ")[-1])
print(f"Document for RAG created with ID: {rag_document_id}")
except (ValueError, IndexError):
pytest.fail("Could not parse document ID from response message.")
try:
chat_payload = {
"prompt": RAG_PROMPT,
"document_id": rag_document_id,
"provider_name": "deepseek",
"load_faiss_retriever": True
}
chat_response = await http_client.post(f"/sessions/{rag_session_id}/chat", json=chat_payload)
# --- MODIFICATION START ---
# If a 500 error occurs, print the detailed response text.
if chat_response.status_code != 200:
print(f"❌ Test Failed! Received status code: {chat_response.status_code}")
print("--- Response Body (for debugging) ---")
print(chat_response.text)
print("---------------------------------------")
assert chat_response.status_code == 200
# --- MODIFICATION END ---
chat_data = chat_response.json()
assert "Jane Doe" in chat_data["answer"]
assert "Nexus" in chat_data["answer"]
print("✅ Chat with document retrieval test passed.")
finally:
# Clean up the document after the test
delete_response = await http_client.delete(f"/documents/{rag_document_id}")
assert delete_response.status_code == 200
print(f"Document {rag_document_id} deleted successfully.")