4acda9c80f
* global exception middleware * redo both logging middlewares as one * remove extra middleware files
96 lines
3.8 KiB
Python
96 lines
3.8 KiB
Python
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from letta.log_context import get_log_context
|
|
from letta.server.rest_api.middleware import LoggingMiddleware
|
|
|
|
|
|
@pytest.fixture
|
|
def app():
|
|
app = FastAPI()
|
|
app.add_middleware(LoggingMiddleware)
|
|
|
|
@app.get("/v1/agents/{agent_id}")
|
|
async def get_agent(agent_id: str):
|
|
context = get_log_context()
|
|
return {"agent_id": agent_id, "context": context}
|
|
|
|
@app.get("/v1/agents/{agent_id}/tools/{tool_id}")
|
|
async def get_agent_tool(agent_id: str, tool_id: str):
|
|
context = get_log_context()
|
|
return {"agent_id": agent_id, "tool_id": tool_id, "context": context}
|
|
|
|
@app.get("/v1/organizations/{org_id}/users/{user_id}")
|
|
async def get_org_user(org_id: str, user_id: str):
|
|
context = get_log_context()
|
|
return {"org_id": org_id, "user_id": user_id, "context": context}
|
|
|
|
return app
|
|
|
|
|
|
@pytest.fixture
|
|
def client(app):
|
|
return TestClient(app)
|
|
|
|
|
|
class TestLogContextMiddleware:
|
|
def test_extracts_actor_id_from_headers(self, client):
|
|
response = client.get("/v1/agents/agent-123e4567-e89b-42d3-8456-426614174000", headers={"user_id": "user-abc123"})
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["context"]["actor_id"] == "user-abc123"
|
|
|
|
def test_extracts_agent_id_from_path(self, client):
|
|
agent_id = "agent-123e4567-e89b-42d3-8456-426614174000"
|
|
response = client.get(f"/v1/agents/{agent_id}")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["context"]["agent_id"] == agent_id
|
|
|
|
def test_extracts_multiple_primitive_ids_from_path(self, client):
|
|
agent_id = "agent-123e4567-e89b-42d3-8456-426614174000"
|
|
tool_id = "tool-987e6543-e21c-42d3-9456-426614174000"
|
|
response = client.get(f"/v1/agents/{agent_id}/tools/{tool_id}")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["context"]["agent_id"] == agent_id
|
|
assert data["context"]["tool_id"] == tool_id
|
|
|
|
def test_extracts_org_id_with_custom_mapping(self, client):
|
|
org_id = "org-123e4567-e89b-42d3-8456-426614174000"
|
|
user_id = "user-987e6543-e21c-42d3-9456-426614174000"
|
|
response = client.get(f"/v1/organizations/{org_id}/users/{user_id}")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["context"]["org_id"] == org_id
|
|
assert data["context"]["user_id"] == user_id
|
|
|
|
def test_extracts_both_header_and_path_context(self, client):
|
|
agent_id = "agent-123e4567-e89b-42d3-8456-426614174000"
|
|
response = client.get(f"/v1/agents/{agent_id}", headers={"user_id": "user-abc123"})
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["context"]["actor_id"] == "user-abc123"
|
|
assert data["context"]["agent_id"] == agent_id
|
|
|
|
def test_handles_request_without_context(self, client):
|
|
response = client.get("/v1/health")
|
|
assert response.status_code == 404
|
|
|
|
def test_context_cleared_between_requests(self, client):
|
|
agent_id_1 = "agent-111e4567-e89b-42d3-8456-426614174000"
|
|
agent_id_2 = "agent-222e4567-e89b-42d3-8456-426614174000"
|
|
|
|
response1 = client.get(f"/v1/agents/{agent_id_1}", headers={"user_id": "user-1"})
|
|
assert response1.status_code == 200
|
|
data1 = response1.json()
|
|
assert data1["context"]["agent_id"] == agent_id_1
|
|
assert data1["context"]["actor_id"] == "user-1"
|
|
|
|
response2 = client.get(f"/v1/agents/{agent_id_2}", headers={"user_id": "user-2"})
|
|
assert response2.status_code == 200
|
|
data2 = response2.json()
|
|
assert data2["context"]["agent_id"] == agent_id_2
|
|
assert data2["context"]["actor_id"] == "user-2"
|