mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
Paul Hernandez
242 lines
7.1 KiB
Python
242 lines
7.1 KiB
Python
"""Tests for MCP tool utilities."""
|
|
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
from httpx import HTTPStatusError, Request
|
|
from mcp.server.fastmcp.exceptions import ToolError
|
|
|
|
from basic_memory.mcp.tools.utils import (
|
|
call_delete,
|
|
call_get,
|
|
call_post,
|
|
call_put,
|
|
get_error_message,
|
|
)
|
|
from basic_memory.workspace_context import (
|
|
WORKSPACE_SLUG_HEADER,
|
|
WORKSPACE_TYPE_HEADER,
|
|
workspace_permalink_context,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_response(monkeypatch):
|
|
"""Create a mock response."""
|
|
|
|
class MockResponse:
|
|
def __init__(self, status_code=200):
|
|
self.status_code = status_code
|
|
self.is_success = status_code < 400
|
|
self.json = lambda: {}
|
|
self.headers = {"content-type": "application/json"}
|
|
|
|
def raise_for_status(self):
|
|
if self.status_code >= 400:
|
|
raise HTTPStatusError(
|
|
message=f"HTTP Error {self.status_code}",
|
|
request=Request("GET", "http://test.com"),
|
|
response=cast(Any, self),
|
|
)
|
|
|
|
return MockResponse
|
|
|
|
|
|
class _Client:
|
|
def __init__(self):
|
|
self.calls: list[tuple[str, tuple, dict]] = []
|
|
self._responses: dict[str, object] = {}
|
|
|
|
def set_response(self, method: str, response):
|
|
self._responses[method.lower()] = response
|
|
|
|
async def get(self, *args, **kwargs):
|
|
self.calls.append(("get", args, kwargs))
|
|
return self._responses["get"]
|
|
|
|
async def post(self, *args, **kwargs):
|
|
self.calls.append(("post", args, kwargs))
|
|
return self._responses["post"]
|
|
|
|
async def put(self, *args, **kwargs):
|
|
self.calls.append(("put", args, kwargs))
|
|
return self._responses["put"]
|
|
|
|
async def delete(self, *args, **kwargs):
|
|
self.calls.append(("delete", args, kwargs))
|
|
return self._responses["delete"]
|
|
|
|
|
|
def _client(client: _Client) -> Any:
|
|
return cast(Any, client)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_get_success(mock_response):
|
|
"""Test successful GET request."""
|
|
client = _Client()
|
|
client.set_response("get", mock_response())
|
|
|
|
response = await call_get(_client(client), "http://test.com")
|
|
assert response.status_code == 200
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_get_error(mock_response):
|
|
"""Test GET request with error."""
|
|
client = _Client()
|
|
client.set_response("get", mock_response(404))
|
|
|
|
with pytest.raises(ToolError) as exc:
|
|
await call_get(_client(client), "http://test.com")
|
|
assert "Resource not found" in str(exc.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_post_success(mock_response):
|
|
"""Test successful POST request."""
|
|
client = _Client()
|
|
response = mock_response()
|
|
response.json = lambda: {"test": "data"}
|
|
client.set_response("post", response)
|
|
|
|
response = await call_post(_client(client), "http://test.com", json={"test": "data"})
|
|
assert response.status_code == 200
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_post_error(mock_response):
|
|
"""Test POST request with error."""
|
|
client = _Client()
|
|
response = mock_response(500)
|
|
response.json = lambda: {"test": "error"}
|
|
|
|
client.set_response("post", response)
|
|
|
|
with pytest.raises(ToolError) as exc:
|
|
await call_post(_client(client), "http://test.com", json={"test": "data"})
|
|
assert "Internal server error" in str(exc.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_put_success(mock_response):
|
|
"""Test successful PUT request."""
|
|
client = _Client()
|
|
client.set_response("put", mock_response())
|
|
|
|
response = await call_put(_client(client), "http://test.com", json={"test": "data"})
|
|
assert response.status_code == 200
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_put_error(mock_response):
|
|
"""Test PUT request with error."""
|
|
client = _Client()
|
|
client.set_response("put", mock_response(400))
|
|
|
|
with pytest.raises(ToolError) as exc:
|
|
await call_put(_client(client), "http://test.com", json={"test": "data"})
|
|
assert "Invalid request" in str(exc.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_delete_success(mock_response):
|
|
"""Test successful DELETE request."""
|
|
client = _Client()
|
|
client.set_response("delete", mock_response())
|
|
|
|
response = await call_delete(_client(client), "http://test.com")
|
|
assert response.status_code == 200
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_delete_error(mock_response):
|
|
"""Test DELETE request with error."""
|
|
client = _Client()
|
|
client.set_response("delete", mock_response(403))
|
|
|
|
with pytest.raises(ToolError) as exc:
|
|
await call_delete(_client(client), "http://test.com")
|
|
assert "Access denied" in str(exc.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_get_with_params(mock_response):
|
|
"""Test GET request with query parameters."""
|
|
client = _Client()
|
|
client.set_response("get", mock_response())
|
|
|
|
params = {"key": "value", "test": "data"}
|
|
await call_get(_client(client), "http://test.com", params=params)
|
|
|
|
assert len(client.calls) == 1
|
|
method, _args, kwargs = client.calls[0]
|
|
assert method == "get"
|
|
assert kwargs["params"] == params
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_post_adds_workspace_permalink_headers_at_request_time(mock_response):
|
|
client = _Client()
|
|
client.set_response("post", mock_response())
|
|
|
|
with workspace_permalink_context("team-paul", "organization"):
|
|
await call_post(
|
|
_client(client),
|
|
"http://test.com",
|
|
headers={"X-Existing": "value"},
|
|
)
|
|
|
|
assert len(client.calls) == 1
|
|
_, _args, kwargs = client.calls[0]
|
|
request_headers = kwargs["headers"]
|
|
assert request_headers["X-Existing"] == "value"
|
|
assert request_headers[WORKSPACE_SLUG_HEADER] == "team-paul"
|
|
assert request_headers[WORKSPACE_TYPE_HEADER] == "organization"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_error_message():
|
|
"""Test the get_error_message function."""
|
|
|
|
# Test 400 status code
|
|
message = get_error_message(400, "http://test.com/resource", "GET")
|
|
assert "Invalid request" in message
|
|
assert "resource" in message
|
|
|
|
# Test 404 status code
|
|
message = get_error_message(404, "http://test.com/missing", "GET")
|
|
assert "Resource not found" in message
|
|
assert "missing" in message
|
|
|
|
# Test 500 status code
|
|
message = get_error_message(500, "http://test.com/server", "POST")
|
|
assert "Internal server error" in message
|
|
assert "server" in message
|
|
|
|
# Test URL object handling
|
|
from httpx import URL
|
|
|
|
url = URL("http://test.com/complex/path")
|
|
message = get_error_message(403, url, "DELETE")
|
|
assert "Access denied" in message
|
|
assert "path" in message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_post_with_json(mock_response):
|
|
"""Test POST request with JSON payload."""
|
|
client = _Client()
|
|
response = mock_response()
|
|
response.json = lambda: {"test": "data"}
|
|
|
|
client.set_response("post", response)
|
|
|
|
json_data = {"key": "value", "nested": {"test": "data"}}
|
|
await call_post(_client(client), "http://test.com", json=json_data)
|
|
|
|
assert len(client.calls) == 1
|
|
method, _args, kwargs = client.calls[0]
|
|
assert method == "post"
|
|
assert kwargs["json"] == json_data
|