import os import json import pytest import uuid from bson.objectid import ObjectId from backend.app import create_app from backend.extensions import mongo # Ensure a valid Gemini API key is provided for integration testing. VALID_GEMINI_KEY = "AIzaSyAMpVRmzQPYAYRH5GiBoQLY-r95ohYmhYs" @pytest.fixture def client(): """ Initializes and yields a Flask test client from create_app(). """ app = create_app() app.config["TESTING"] = True with app.test_client() as test_client: yield test_client @pytest.fixture def auth_headers(client): """ Registers and logs in a new test user. Returns a dictionary containing the Authorization header and the user_id. """ unique_suffix = str(uuid.uuid4())[:8] username = f"test_dialog_{unique_suffix}" email = f"{username}@example.com" reg_payload = { "username": username, "email": email, "password": "Password123" } reg_resp = client.post( "/api/register", data=json.dumps(reg_payload), content_type="application/json" ) assert reg_resp.status_code == 201, f"User registration failed: {reg_resp.data.decode()}" reg_data = reg_resp.get_json() return {"Authorization": f"Bearer {reg_data['token']}", "user_id": reg_data["user_id"]} @pytest.fixture def gemini_api_config(auth_headers): """ Inserts a Gemini API key document for the test user into the api_list collection. Uses the valid API key from environment variable. """ from datetime import datetime # Import the datetime class for utcnow() user_id = ObjectId(auth_headers["user_id"]) api_payload = { "uid": user_id, "name": "Gemini", "key": VALID_GEMINI_KEY, "selected": True, "createdAt": datetime.utcnow(), "updatedAt": datetime.utcnow() } inserted_result = mongo.db.api_list.insert_one(api_payload) yield # Teardown: remove the API key document after test. mongo.db.api_list.delete_one({"_id": inserted_result.inserted_id}) @pytest.fixture def project_id(client, auth_headers): """ Creates a new project for the test user and returns its ID. """ payload = { "name": "Dialog Test Project", "topic": "Integration Testing", "description": "Project created for testing dialog endpoints." } resp = client.post("/api/projects", data=json.dumps(payload), content_type="application/json", headers={"Authorization": auth_headers["Authorization"]}) assert resp.status_code == 201, f"Project creation failed: {resp.data.decode()}" data = resp.get_json() return data["project_id"] def test_create_dialog_no_start_message(client, auth_headers, gemini_api_config, project_id): """ Test creating a dialog session without a start message. """ payload = {"projectId": project_id} resp = client.post("/api/dialog", data=json.dumps(payload), headers={"Authorization": auth_headers["Authorization"]}, content_type="application/json") assert resp.status_code == 201, f"Expected 201, got {resp.status_code}, {resp.data.decode()}" data = resp.get_json() assert "dialog_id" in data # Retrieve the new dialog session to verify dialog_id = data["dialog_id"] get_resp = client.get(f"/api/dialog/{dialog_id}", headers={"Authorization": auth_headers["Authorization"]}) assert get_resp.status_code == 200, f"Failed to get dialog: {get_resp.data.decode()}" dialog_data = get_resp.get_json() assert "sessionStartedAt" in dialog_data def test_create_dialog_with_start_message(client, auth_headers, gemini_api_config, project_id): """ Test creating a dialog session with a start message. """ payload = { "projectId": project_id, "sessionId": "testSession123", "startMessage": "Hello, I need research guidance." } resp = client.post("/api/dialog", data=json.dumps(payload), headers={"Authorization": auth_headers["Authorization"]}, content_type="application/json") assert resp.status_code == 201, f"Expected 201, got {resp.status_code}, {resp.data.decode()}" data = resp.get_json() dialog_id = data["dialog_id"] get_resp = client.get(f"/api/dialog/{dialog_id}", headers={"Authorization": auth_headers["Authorization"]}) assert get_resp.status_code == 200, f"Failed to retrieve dialog: {get_resp.data.decode()}" dialog_data = get_resp.get_json() msgs = dialog_data.get("messages", []) assert len(msgs) >= 1, "Expected at least one message in the dialog." assert msgs[0]["role"] == "user" assert "Hello, I need research guidance." in msgs[0]["content"] def test_list_dialogs(client, auth_headers, gemini_api_config, project_id): """ Creates a couple of dialogs and then lists dialogs filtered by project. """ for _ in range(2): payload = {"projectId": project_id} resp = client.post("/api/dialog", data=json.dumps(payload), headers={"Authorization": auth_headers["Authorization"]}, content_type="application/json") assert resp.status_code == 201 list_resp = client.get(f"/api/dialog?projectId={project_id}", headers={"Authorization": auth_headers["Authorization"]}) assert list_resp.status_code == 200, f"Listing dialogs failed: {list_resp.data.decode()}" data = list_resp.get_json() dialogs = data.get("dialogs", []) assert len(dialogs) >= 2, "Expected at least two dialog sessions." def test_send_dialog_message_real_gemini(client, auth_headers, gemini_api_config, project_id): """ Test sending a message in a dialog session using the vector-based prompt. This test interacts with the real Gemini API. """ # Create a new dialog session. create_payload = {"projectId": project_id} create_resp = client.post("/api/dialog", data=json.dumps(create_payload), headers={"Authorization": auth_headers["Authorization"]}, content_type="application/json") assert create_resp.status_code == 201, f"Dialog creation failed: {create_resp.data.decode()}" dialog_id = create_resp.get_json()["dialog_id"] # Send a message. send_payload = {"content": "What further research should I pursue based on my current websites?"} send_resp = client.post(f"/api/dialog/{dialog_id}/send", data=json.dumps(send_payload), headers={"Authorization": auth_headers["Authorization"]}, content_type="application/json") # This test makes a live call to the Gemini API. assert send_resp.status_code == 200, f"Send message failed: {send_resp.data.decode()}" send_data = send_resp.get_json() assert "llmResponse" in send_data, "Response missing LLM response." print("Gemini LLM response:", send_data["llmResponse"]) # Verify that the dialog now has additional messages. get_resp = client.get(f"/api/dialog/{dialog_id}", headers={"Authorization": auth_headers["Authorization"]}) assert get_resp.status_code == 200, f"Retrieving dialog failed: {get_resp.data.decode()}" dialog_data = get_resp.get_json() messages = dialog_data.get("messages", []) assert len(messages) >= 2, "Expected at least two messages after sending (user and system)." def test_end_and_delete_session(client, auth_headers, gemini_api_config, project_id): """ Tests ending a dialog session and then deleting it. """ create_payload = {"projectId": project_id} create_resp = client.post("/api/dialog", data=json.dumps(create_payload), headers={"Authorization": auth_headers["Authorization"]}, content_type="application/json") assert create_resp.status_code == 201, f"Dialog creation failed: {create_resp.data.decode()}" dialog_id = create_resp.get_json()["dialog_id"] # End the dialog session. end_resp = client.put(f"/api/dialog/{dialog_id}/end", headers={"Authorization": auth_headers["Authorization"]}) assert end_resp.status_code == 200, f"Ending dialog failed: {end_resp.data.decode()}" # Attempt to send another message; should fail. send_payload = {"content": "Trying to send after end."} send_resp = client.post(f"/api/dialog/{dialog_id}/send", data=json.dumps(send_payload), headers={"Authorization": auth_headers["Authorization"]}, content_type="application/json") assert send_resp.status_code == 400, "Expected error when sending message after ending session." # Delete the dialog session. del_resp = client.delete(f"/api/dialog/{dialog_id}", headers={"Authorization": auth_headers["Authorization"]}) assert del_resp.status_code == 200, f"Deleting dialog failed: {del_resp.data.decode()}" # Verify that retrieving the dialog now returns 404. get_resp = client.get(f"/api/dialog/{dialog_id}", headers={"Authorization": auth_headers["Authorization"]}) assert get_resp.status_code == 404, "Expected 404 when retrieving a deleted dialog."