import pytest import httpx # The base URL for the local server BASE_URL = "http://127.0.0.1:8000" TEST_PROMPT = "Explain the theory of relativity in one sentence." # Global variables to pass state between sequential tests created_document_id = None created_session_id = None async def test_root_endpoint(): """Tests if the root endpoint is alive.""" print("\n--- Running test_root_endpoint ---") async with httpx.AsyncClient() as client: response = await client.get(f"{BASE_URL}/") assert response.status_code == 200 assert response.json() == {"status": "AI Model Hub is running!"} print("✅ Root endpoint test passed.") # --- Session and Chat Lifecycle Tests --- async def test_create_session(): """ Tests creating a new chat session and saves the ID for the next test. """ global created_session_id print("\n--- Running test_create_session ---") url = f"{BASE_URL}/sessions" payload = {"user_id": "integration_tester", "model": "deepseek"} async with httpx.AsyncClient() as client: response = await client.post(url, json=payload) assert response.status_code == 200, f"Failed to create session. Response: {response.text}" response_data = response.json() assert "id" in response_data assert response_data["user_id"] == "integration_tester" assert response_data["model_name"] == "deepseek" created_session_id = response_data["id"] print(f"✅ Session created successfully with ID: {created_session_id}") async def test_chat_in_session(): """ Tests sending a message within the session created by the previous test. """ print("\n--- Running test_chat_in_session ---") assert created_session_id is not None, "Session ID was not set by the create_session test." url = f"{BASE_URL}/sessions/{created_session_id}/chat" payload = {"prompt": TEST_PROMPT} async with httpx.AsyncClient(timeout=60.0) as client: response = await client.post(url, json=payload) assert response.status_code == 200, f"Chat request failed. Response: {response.text}" response_data = response.json() assert "answer" in response_data assert len(response_data["answer"]) > 0 assert response_data["model_used"] == "deepseek" print("✅ Chat in session test passed.") # --- Document Management Lifecycle Tests --- async def test_add_document_for_lifecycle(): """ Adds a document and saves its ID to be used by the list and delete tests. """ global created_document_id print("\n--- Running test_add_document (for lifecycle) ---") url = f"{BASE_URL}/documents" doc_data = {"title": "Lifecycle Test Doc", "text": "This doc will be listed and deleted."} async with httpx.AsyncClient(timeout=30.0) as client: response = await client.post(url, json=doc_data) assert response.status_code == 200, f"Failed to add document. Response: {response.text}" response_data = response.json() message = response_data.get("message", "") assert "added successfully with ID" in message try: created_document_id = int(message.split(" with ID ")[-1]) except (ValueError, IndexError): pytest.fail("Could not parse document ID from response message.") print(f"✅ Document for lifecycle test created with ID: {created_document_id}") async def test_list_documents(): """ Tests listing documents to ensure the previously created one appears. """ print("\n--- Running test_list_documents ---") assert created_document_id is not None, "Document ID was not set by the add test." url = f"{BASE_URL}/documents" async with httpx.AsyncClient() as client: response = await client.get(url) assert response.status_code == 200 response_data = response.json() assert "documents" in response_data ids_in_response = {doc["id"] for doc in response_data["documents"]} assert created_document_id in ids_in_response print("✅ Document list test passed.") async def test_delete_document(): """ Tests deleting the document created at the start of the lifecycle. """ print("\n--- Running test_delete_document ---") assert created_document_id is not None, "Document ID was not set by the add test." url = f"{BASE_URL}/documents/{created_document_id}" async with httpx.AsyncClient() as client: response = await client.delete(url) assert response.status_code == 200 response_data = response.json() assert response_data["message"] == "Document deleted successfully" assert response_data["document_id"] == created_document_id print("✅ Document delete test passed.")