219 lines
9.2 KiB
Python
219 lines
9.2 KiB
Python
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."
|