mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
phernandez
250 lines
7.9 KiB
Python
250 lines
7.9 KiB
Python
"""Tests for knowledge graph API endpoints."""
|
|
from typing import AsyncGenerator
|
|
import pytest
|
|
import pytest_asyncio
|
|
from fastapi import FastAPI
|
|
from httpx import AsyncClient, ASGITransport
|
|
from icecream import ic
|
|
from loguru import logger
|
|
|
|
from basic_memory.deps import get_project_config, get_engine
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def app(test_config, engine) -> FastAPI:
|
|
"""Create FastAPI test application."""
|
|
# Lazy import router to avoid app startup issues
|
|
from basic_memory.api.routers.knowledge import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
|
|
app.dependency_overrides[get_project_config] = lambda: test_config
|
|
app.dependency_overrides[get_engine] = lambda: engine
|
|
return app
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def client(app: FastAPI) -> AsyncGenerator[AsyncClient, None]:
|
|
"""Create client using ASGI transport - same as CLI will use."""
|
|
async with AsyncClient(
|
|
transport=ASGITransport(app=app),
|
|
base_url="http://test"
|
|
) as client:
|
|
yield client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_entities(client: AsyncClient):
|
|
"""Should create entities successfully."""
|
|
# Create an entity
|
|
response = await client.post("/knowledge/entities", json={
|
|
"entities": [{
|
|
"name": "Test Entity",
|
|
"entity_type": "test"
|
|
}]
|
|
})
|
|
|
|
# Verify creation
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert len(data["entities"]) == 1
|
|
entity = data["entities"][0]
|
|
assert entity["id"] == "test/test_entity"
|
|
assert entity["name"] == "Test Entity"
|
|
|
|
entity_type = entity.get("entity_type") or entity.get("entityType")
|
|
assert entity_type == "test"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_entity(client: AsyncClient):
|
|
"""Should retrieve an entity by ID."""
|
|
# First create an entity
|
|
create_response = await client.post("/knowledge/entities", json={
|
|
"entities": [{
|
|
"name": "Test Entity",
|
|
"entity_type": "test",
|
|
}]
|
|
})
|
|
entity_id = create_response.json()["entities"][0]["id"]
|
|
|
|
# Now get it by ID
|
|
response = await client.get(f"/knowledge/entities/{entity_id}")
|
|
|
|
# Verify retrieval
|
|
assert response.status_code == 200
|
|
entity = response.json()
|
|
assert entity["id"] == entity_id
|
|
assert entity["name"] == "Test Entity"
|
|
|
|
entity_type = entity.get("entity_type") or entity.get("entityType")
|
|
assert entity_type == "test"
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_relations(client: AsyncClient):
|
|
"""Should create relations between entities."""
|
|
# Create two entities to relate
|
|
entities = [
|
|
{"name": "Source Entity", "entity_type": "test"},
|
|
{"name": "Target Entity", "entity_type": "test"}
|
|
]
|
|
create_response = await client.post("/knowledge/entities", json={"entities": entities})
|
|
created = create_response.json()["entities"]
|
|
source_id = created[0]["id"]
|
|
target_id = created[1]["id"]
|
|
|
|
# Create relation between them
|
|
response = await client.post("/knowledge/relations", json={
|
|
"relations": [{
|
|
"from_id": source_id,
|
|
"to_id": target_id,
|
|
"relation_type": "related_to"
|
|
}]
|
|
})
|
|
|
|
# Verify relation
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert len(data["relations"]) == 1
|
|
relation = data["relations"][0]
|
|
assert relation["from_id"] == source_id
|
|
assert relation["to_id"] == target_id
|
|
assert relation["relation_type"] == "related_to"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_observations(client: AsyncClient):
|
|
"""Should add observations to an entity."""
|
|
# Create an entity first
|
|
create_response = await client.post("/knowledge/entities", json={
|
|
"entities": [{
|
|
"name": "Test Entity",
|
|
"entity_type": "test"
|
|
}]
|
|
})
|
|
entity_id = create_response.json()["entities"][0]["id"]
|
|
|
|
# Add observations
|
|
response = await client.post("/knowledge/observations", json={
|
|
"entity_id": entity_id,
|
|
"observations": [
|
|
"First observation",
|
|
"Second observation"
|
|
]
|
|
})
|
|
|
|
logger.debug(ic(response.content))
|
|
|
|
# Verify observations were added
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["entity_id"] == entity_id
|
|
assert len(data["observations"]) == 2
|
|
assert data["observations"][0]["content"] == "First observation"
|
|
assert data["observations"][1]["content"] == "Second observation"
|
|
|
|
# Verify observations appear in entity
|
|
entity_response = await client.get(f"/knowledge/entities/{entity_id}")
|
|
entity = entity_response.json()
|
|
assert len(entity["observations"]) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_nodes(client: AsyncClient):
|
|
"""Should search for entities in the knowledge graph."""
|
|
# Create a few entities with different names
|
|
entities = [
|
|
{"name": "Not found", "entity_type": "negative"},
|
|
{"name": "Alpha Test", "entity_type": "test"},
|
|
{"name": "Beta Test", "entity_type": "test"},
|
|
{"name": "Gamma Production", "entity_type": "test"} #match entity_type
|
|
]
|
|
await client.post("/knowledge/entities", json={"entities": entities})
|
|
|
|
# Search for "Test" in names
|
|
response = await client.post("/knowledge/search", json={"query": "Test"})
|
|
|
|
# Verify search results
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["query"] == "Test"
|
|
assert len(data["matches"]) == 3
|
|
names = [entity["name"] for entity in data["matches"]]
|
|
assert "Alpha Test" in names
|
|
assert "Beta Test" in names
|
|
assert "Gamma Production" in names
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_knowledge_flow(client: AsyncClient):
|
|
"""Test a complete knowledge graph flow with multiple operations."""
|
|
# 1. Create main entity
|
|
main_response = await client.post("/knowledge/entities", json={
|
|
"entities": [{
|
|
"name": "Main Entity",
|
|
"entity_type": "test"
|
|
},
|
|
{
|
|
"name": "Non Entity",
|
|
"entity_type": "n_a"
|
|
}]
|
|
})
|
|
main_id = main_response.json()["entities"][0]["id"]
|
|
assert main_response.status_code == 200
|
|
assert main_id is not None
|
|
|
|
# 2. Create related entities
|
|
related_response = await client.post("/knowledge/entities", json={
|
|
"entities": [
|
|
{"name": "Related One", "entity_type": "test"},
|
|
{"name": "Related Two", "entity_type": "test"}
|
|
]
|
|
})
|
|
related = related_response.json()["entities"]
|
|
related_ids = [e["id"] for e in related]
|
|
assert related_response.status_code == 200
|
|
assert len(related_ids) == 2
|
|
|
|
# 3. Add relations
|
|
relations_response = await client.post("/knowledge/relations", json={
|
|
"relations": [
|
|
{
|
|
"from_id": main_id,
|
|
"to_id": related_ids[0],
|
|
"relation_type": "connects_to"
|
|
},
|
|
{
|
|
"from_id": main_id,
|
|
"to_id": related_ids[1],
|
|
"relation_type": "connects_to"
|
|
}
|
|
]
|
|
})
|
|
assert relations_response.status_code == 200
|
|
assert len(relations_response.json()["relations"]) == 2
|
|
|
|
# 4. Add observations to main entity
|
|
await client.post("/knowledge/observations", json={
|
|
"entity_id": main_id,
|
|
"observations": [
|
|
"Connected to first related entity",
|
|
"Connected to second related entity"
|
|
]
|
|
})
|
|
|
|
# 5. Verify full graph structure
|
|
main_get = await client.get(f"/knowledge/entities/{main_id}")
|
|
main_entity = main_get.json()
|
|
|
|
# Check entity structure
|
|
assert main_entity["name"] == "Main Entity"
|
|
assert len(main_entity["observations"]) == 2
|
|
assert len(main_entity["relations"]) == 2
|
|
|
|
# 6. Search should find all related entities
|
|
search = await client.post("/knowledge/search", json={"query": "Related"})
|
|
matches = search.json()["matches"]
|
|
assert len(matches) == 3 # Should find both related entities |