mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
fix tests
This commit is contained in:
phernandez
@@ -127,15 +127,7 @@ async def delete_relations(
|
||||
data: DeleteRelationsRequest, knowledge_service: KnowledgeServiceDep
|
||||
) -> CreateEntityResponse:
|
||||
"""Delete relations between entities."""
|
||||
to_delete = [
|
||||
{
|
||||
"from_id": relation.from_id,
|
||||
"to_id": relation.to_id,
|
||||
"relation_type": relation.relation_type,
|
||||
}
|
||||
for relation in data.relations
|
||||
]
|
||||
updated_entities = await knowledge_service.delete_relations(to_delete)
|
||||
updated_entities = await knowledge_service.delete_relations(data.relations)
|
||||
return CreateEntityResponse(
|
||||
entities=[EntityResponse.model_validate(entity) for entity in updated_entities]
|
||||
)
|
||||
|
||||
@@ -79,8 +79,12 @@ class EntityRepository(Repository[Entity]):
|
||||
def get_load_options(self) -> List[LoaderOption]:
|
||||
return [
|
||||
selectinload(Entity.observations),
|
||||
# Load from_relations and both entities for each relation
|
||||
selectinload(Entity.from_relations).selectinload(Relation.from_entity),
|
||||
selectinload(Entity.from_relations).selectinload(Relation.to_entity),
|
||||
# Load to_relations and both entities for each relation
|
||||
selectinload(Entity.to_relations).selectinload(Relation.from_entity),
|
||||
selectinload(Entity.to_relations).selectinload(Relation.to_entity),
|
||||
]
|
||||
|
||||
async def find_by_path_ids(self, path_ids: List[str]) -> Sequence[Entity]:
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
"""Repository for managing Relation objects."""
|
||||
|
||||
from typing import Sequence, List
|
||||
from sqlalchemy import and_
|
||||
from typing import Sequence, List, Optional
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker
|
||||
from sqlalchemy.orm import selectinload
|
||||
from sqlalchemy.orm import selectinload, aliased
|
||||
from sqlalchemy.orm.interfaces import LoaderOption
|
||||
|
||||
from basic_memory.models import Relation, Entity
|
||||
@@ -17,6 +17,25 @@ class RelationRepository(Repository[Relation]):
|
||||
def __init__(self, session_maker: async_sessionmaker):
|
||||
super().__init__(session_maker, Relation)
|
||||
|
||||
async def find_relation(self, from_path_id: str, to_path_id: str, relation_type: str) -> Optional[Relation]:
|
||||
"""Find a relation by its from and to path IDs."""
|
||||
from_entity = aliased(Entity)
|
||||
to_entity = aliased(Entity)
|
||||
|
||||
query = (
|
||||
select(Relation)
|
||||
.join(from_entity, Relation.from_id == from_entity.id)
|
||||
.join(to_entity, Relation.to_id == to_entity.id)
|
||||
.where(
|
||||
and_(
|
||||
from_entity.path_id == from_path_id,
|
||||
to_entity.path_id == to_path_id,
|
||||
Relation.relation_type == relation_type
|
||||
)
|
||||
)
|
||||
)
|
||||
return await self.find_one(query)
|
||||
|
||||
async def find_by_entity(self, from_entity_id: int) -> Sequence[Relation]:
|
||||
"""Find all relations from a specific entity."""
|
||||
query = select(Relation).filter(Relation.from_id == from_entity_id)
|
||||
|
||||
@@ -8,3 +8,8 @@ class EntityNotFoundError(Exception):
|
||||
"""Raised when an entity cannot be found"""
|
||||
|
||||
pass
|
||||
|
||||
class EntityCreationError(Exception):
|
||||
"""Raised when an entity cannot be created"""
|
||||
|
||||
pass
|
||||
|
||||
@@ -7,6 +7,7 @@ from loguru import logger
|
||||
from basic_memory.models import Entity as EntityModel
|
||||
from basic_memory.schemas import Entity as EntitySchema
|
||||
from .files import FileOperations
|
||||
from ..exceptions import EntityCreationError, EntityNotFoundError
|
||||
|
||||
|
||||
class EntityOperations(FileOperations):
|
||||
@@ -49,7 +50,7 @@ class EntityOperations(FileOperations):
|
||||
for entity in entities:
|
||||
created_entity = await self.create_entity(entity)
|
||||
created.append(created_entity)
|
||||
|
||||
|
||||
return created
|
||||
|
||||
async def delete_entity(self, path_id: str) -> bool:
|
||||
@@ -59,8 +60,6 @@ class EntityOperations(FileOperations):
|
||||
try:
|
||||
# Get entity first for file deletion
|
||||
entity = await self.entity_service.get_by_path_id(path_id)
|
||||
if not entity:
|
||||
return True # Already deleted
|
||||
|
||||
# Delete file first (it's source of truth)
|
||||
path = self.get_entity_path(entity)
|
||||
@@ -68,7 +67,11 @@ class EntityOperations(FileOperations):
|
||||
|
||||
# Delete from DB (this will cascade to observations/relations)
|
||||
return await self.entity_service.delete_entity(path_id)
|
||||
|
||||
|
||||
except EntityNotFoundError:
|
||||
logger.info(f"Entity not found: {path_id}")
|
||||
return True # Already deleted
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete entity: {e}")
|
||||
raise
|
||||
|
||||
@@ -36,10 +36,11 @@ class ObservationOperations(RelationOperations):
|
||||
updated_entity = await self.entity_service.get_by_path_id(path_id)
|
||||
|
||||
# Write updated file and checksum
|
||||
checksum = await self.write_entity_file(entity)
|
||||
_, checksum = await self.write_entity_file(entity)
|
||||
await self.entity_service.update_entity(path_id, {"checksum": checksum})
|
||||
|
||||
return updated_entity
|
||||
# query to fetch all relations
|
||||
return await self.entity_service.get_by_path_id(path_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to add observations: {e}")
|
||||
@@ -59,7 +60,7 @@ class ObservationOperations(RelationOperations):
|
||||
await self.observation_service.delete_observations(entity.id, observations)
|
||||
|
||||
# Write updated file
|
||||
checksum = await self.write_entity_file(entity)
|
||||
_, checksum = await self.write_entity_file(entity)
|
||||
await self.entity_service.update_entity(path_id, {"checksum": checksum})
|
||||
|
||||
# Get final entity with all updates
|
||||
|
||||
@@ -66,20 +66,26 @@ class RelationOperations(EntityOperations):
|
||||
# select again to eagerly load all relations
|
||||
return await self.entity_service.open_nodes([e.path_id for e in updated_entities])
|
||||
|
||||
async def delete_relations(self, to_delete: List[Dict[str, Any]]) -> Sequence[EntityModel]:
|
||||
async def delete_relations(self, to_delete: List[RelationSchema]) -> Sequence[EntityModel]:
|
||||
"""Delete relations and return all updated entities."""
|
||||
logger.debug(f"Deleting {len(to_delete)} relations")
|
||||
updated_entities = []
|
||||
entities_to_update = set()
|
||||
|
||||
relations = []
|
||||
|
||||
try:
|
||||
# Delete relations from DB
|
||||
for relation in to_delete:
|
||||
entities_to_update.add(relation["from_id"])
|
||||
entities_to_update.add(relation["to_id"])
|
||||
entities_to_update.add(relation.from_id)
|
||||
entities_to_update.add(relation.to_id)
|
||||
|
||||
relation = await self.relation_service.find_relation(relation.from_id, relation.to_id, relation.relation_type)
|
||||
if relation:
|
||||
relations.append(relation)
|
||||
|
||||
deleted = await self.relation_service.delete_relations(to_delete)
|
||||
if not deleted:
|
||||
# pass Relation models to delete
|
||||
num_deleted = await self.relation_service.delete_relations(relations)
|
||||
if num_deleted == 0:
|
||||
logger.warning("No relations were deleted")
|
||||
|
||||
# Get fresh copies of all updated entities
|
||||
@@ -91,7 +97,7 @@ class RelationOperations(EntityOperations):
|
||||
raise EntityNotFoundError(f"Entity not found: {path_id}")
|
||||
|
||||
# Write updated file
|
||||
checksum = await self.write_entity_file(entity)
|
||||
_, checksum = await self.write_entity_file(entity)
|
||||
updated = await self.entity_service.update_entity(
|
||||
path_id, {"checksum": checksum}
|
||||
)
|
||||
|
||||
@@ -23,6 +23,9 @@ class RelationService(BaseService[RelationRepository]):
|
||||
logger.debug(f"Creating relation: {relation}")
|
||||
return await self.repository.add(relation)
|
||||
|
||||
async def find_relation(self, from_path_id: str, to_path_id: str, relation_type: str) -> Relation:
|
||||
return await self.repository.find_relation(from_path_id, to_path_id, relation_type)
|
||||
|
||||
async def delete_relation(
|
||||
self, from_entity: Entity, to_entity: Entity, relation_type: str
|
||||
) -> bool:
|
||||
@@ -37,20 +40,13 @@ class RelationService(BaseService[RelationRepository]):
|
||||
)
|
||||
return result
|
||||
|
||||
async def delete_relations(self, relations: List[Dict[str, Any]]) -> bool:
|
||||
async def delete_relations(self, relations: List[Relation]) -> int:
|
||||
"""Delete relations matching specified criteria."""
|
||||
logger.debug(f"Deleting {len(relations)} relations")
|
||||
deleted = False
|
||||
for relation in relations:
|
||||
filters = {"from_id": relation["from_id"], "to_id": relation["to_id"]}
|
||||
if "relation_type" in relation:
|
||||
filters["relation_type"] = relation["relation_type"]
|
||||
|
||||
result = await self.repository.delete_by_fields(**filters)
|
||||
if result:
|
||||
deleted = True
|
||||
|
||||
return deleted
|
||||
|
||||
ids = [relation.id for relation in relations]
|
||||
return await self.repository.delete_by_ids(ids)
|
||||
|
||||
async def create_relations(self, relations: List[Relation]) -> Sequence[Relation]:
|
||||
"""Create multiple relations between entities."""
|
||||
|
||||
@@ -308,10 +308,10 @@ async def test_delete_relations(client, relation_repository):
|
||||
async def test_delete_nonexistent_entity(client: AsyncClient):
|
||||
"""Test deleting a nonexistent entity by path ID."""
|
||||
response = await client.post(
|
||||
"/knowledge/entities/delete", json={"entity_ids": ["test/NonExistent"]}
|
||||
"/knowledge/entities/delete", json={"entity_ids": ["test/non_existent"]}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"deleted": False}
|
||||
assert response.json() == {"deleted": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -336,8 +336,8 @@ async def test_delete_nonexistent_relations(client: AsyncClient):
|
||||
request_data = {
|
||||
"relations": [
|
||||
{
|
||||
"from_id": "test/NonExistent1",
|
||||
"to_id": "test/NonExistent2",
|
||||
"from_id": "test/non_existent1",
|
||||
"to_id": "test/non_existent2",
|
||||
"relation_type": "nonexistent",
|
||||
}
|
||||
]
|
||||
@@ -389,13 +389,13 @@ async def test_full_knowledge_flow(client: AsyncClient):
|
||||
json={
|
||||
"relations": [
|
||||
{
|
||||
"from_id": "test/MainEntity",
|
||||
"to_id": "test/RelatedOne",
|
||||
"from_id": "test/main_entity",
|
||||
"to_id": "test/related_one",
|
||||
"relation_type": "connects_to",
|
||||
},
|
||||
{
|
||||
"from_id": "test/MainEntity",
|
||||
"to_id": "test/RelatedTwo",
|
||||
"from_id": "test/main_entity",
|
||||
"to_id": "test/related_two",
|
||||
"relation_type": "connects_to",
|
||||
},
|
||||
]
|
||||
@@ -407,7 +407,7 @@ async def test_full_knowledge_flow(client: AsyncClient):
|
||||
await client.post(
|
||||
"/knowledge/observations",
|
||||
json={
|
||||
"entity_id": "test/MainEntity",
|
||||
"entity_id": "test/main_entity",
|
||||
"observations": [
|
||||
"Connected to first related entity",
|
||||
"Connected to second related entity",
|
||||
@@ -428,7 +428,7 @@ async def test_full_knowledge_flow(client: AsyncClient):
|
||||
# 6. Search should find all related entities
|
||||
search = await client.post("/knowledge/search", json={"query": "Related"})
|
||||
matches = search.json()["matches"]
|
||||
assert len(matches) == 2 # Should find both related entities
|
||||
assert len(matches) == 3 # Should find both related entities, and the main one with the observation
|
||||
|
||||
# 7. Delete main entity
|
||||
response = await client.post(
|
||||
|
||||
@@ -87,7 +87,7 @@ async def test_delete_observations(session_maker: async_sessionmaker, repo):
|
||||
"""Test deleting observations by entity_id."""
|
||||
# Create test entity
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
entity = Entity(name="test_entity", entity_type="test", description="Test entity")
|
||||
entity = Entity(name="test_entity", entity_type="test", description="Test entity", path_id="test/test_entity")
|
||||
session.add(entity)
|
||||
await session.flush()
|
||||
|
||||
@@ -110,7 +110,7 @@ async def test_delete_observation_by_id(session_maker: async_sessionmaker, repo)
|
||||
"""Test deleting a single observation by its ID."""
|
||||
# Create test entity
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
entity = Entity(name="test_entity", entity_type="test", description="Test entity")
|
||||
entity = Entity(name="test_entity", entity_type="test", description="Test entity", path_id="test/test_entity")
|
||||
session.add(entity)
|
||||
await session.flush()
|
||||
|
||||
@@ -132,7 +132,7 @@ async def test_delete_observation_by_content(session_maker: async_sessionmaker,
|
||||
"""Test deleting observations by content."""
|
||||
# Create test entity
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
entity = Entity(name="test_entity", entity_type="test", description="Test entity")
|
||||
entity = Entity(name="test_entity", entity_type="test", description="Test entity", path_id="test/test_entity")
|
||||
session.add(entity)
|
||||
await session.flush()
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@ async def source_entity(session_maker):
|
||||
entity = Entity(
|
||||
name="test_source",
|
||||
entity_type="source",
|
||||
path_id="source/test_source",
|
||||
description="Source entity",
|
||||
)
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
@@ -29,6 +30,7 @@ async def target_entity(session_maker):
|
||||
entity = Entity(
|
||||
name="test_target",
|
||||
entity_type="target",
|
||||
path_id="target/test_target",
|
||||
description="Target entity",
|
||||
)
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
@@ -56,6 +58,7 @@ async def related_entity(entity_repository):
|
||||
entity_data = {
|
||||
"name": "Related Entity",
|
||||
"entity_type": "test",
|
||||
"path_id": "test/related_entity",
|
||||
"description": "A related test entity",
|
||||
"references": "",
|
||||
}
|
||||
@@ -151,6 +154,14 @@ async def test_find_by_entities(
|
||||
assert relations[0].id == sample_relation.id
|
||||
assert relations[0].relation_type == sample_relation.relation_type
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_relation(relation_repository: RelationRepository, sample_relation: Relation):
|
||||
"""Test finding relations by type"""
|
||||
relation = await relation_repository.find_relation(from_path_id=sample_relation.from_entity.path_id,
|
||||
to_path_id=sample_relation.to_entity.path_id,
|
||||
relation_type=sample_relation.relation_type)
|
||||
assert relation.id == sample_relation.id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_by_type(relation_repository: RelationRepository, sample_relation: Relation):
|
||||
|
||||
@@ -100,20 +100,19 @@ def test_entity_out_from_attributes():
|
||||
"""Test EntityOut creation from database model attributes."""
|
||||
# Simulate database model attributes
|
||||
db_data = {
|
||||
"id": "123",
|
||||
"path_id": "test/test",
|
||||
"name": "test",
|
||||
"entity_type": "test",
|
||||
"description": "test description",
|
||||
"observations": [{"id": 1, "content": "test obs", "context": None}],
|
||||
"relations": [
|
||||
{"id": 1, "from_id": 123, "to_id": 456, "relation_type": "test", "context": None}
|
||||
{"id": 1, "from_id": "test/test", "to_id": "test/test", "relation_type": "test", "context": None}
|
||||
],
|
||||
}
|
||||
entity = EntityResponse.model_validate(db_data)
|
||||
assert entity.id == 123
|
||||
assert entity.path_id == "test/test"
|
||||
assert entity.description == "test description"
|
||||
assert len(entity.observations) == 1
|
||||
assert entity.observations[0].id == 1
|
||||
assert len(entity.relations) == 1
|
||||
|
||||
|
||||
@@ -155,7 +154,7 @@ def test_search_nodes_input():
|
||||
|
||||
def test_open_nodes_input():
|
||||
"""Test OpenNodesInput validation."""
|
||||
open_input = OpenNodesRequest.model_validate({"entity_ids": [1, 2]})
|
||||
open_input = OpenNodesRequest.model_validate({"entity_ids": ["test", "test2"]})
|
||||
assert len(open_input.entity_ids) == 2
|
||||
|
||||
# Empty names list should fail
|
||||
|
||||
@@ -15,9 +15,10 @@ async def file_sync_service(document_repository) -> FileSyncService:
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def docs_dir(test_project_path) -> Path:
|
||||
async def docs_dir(test_config) -> Path:
|
||||
"""Get documents directory."""
|
||||
return test_project_path / "documents"
|
||||
test_config.documents_dir.mkdir(parents=True)
|
||||
return test_config.documents_dir
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
@@ -26,8 +27,8 @@ async def sample_files(docs_dir) -> dict[str, str]:
|
||||
# Create test structure
|
||||
design_dir = docs_dir / "design"
|
||||
notes_dir = docs_dir / "notes"
|
||||
design_dir.mkdir(exist_ok=True)
|
||||
notes_dir.mkdir(exist_ok=True)
|
||||
design_dir.mkdir(parents=True, exist_ok=True)
|
||||
notes_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Map of relative paths to content
|
||||
files = {
|
||||
|
||||
@@ -4,11 +4,12 @@ from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from basic_memory.models import Entity as EntityModel
|
||||
from basic_memory.schemas import Entity as EntitySchema, Relation as RelationSchema
|
||||
from basic_memory.services import EntityService
|
||||
from basic_memory.services.exceptions import EntityNotFoundError, FileOperationError
|
||||
from basic_memory.services.exceptions import EntityNotFoundError, FileOperationError, EntityCreationError
|
||||
from basic_memory.services.knowledge import KnowledgeService
|
||||
|
||||
|
||||
@@ -17,7 +18,7 @@ async def test_get_entity_path(knowledge_service: KnowledgeService):
|
||||
"""Should generate correct filesystem path for entity."""
|
||||
entity = EntityModel(id=1, name="test-entity", entity_type="concept", description="Test entity")
|
||||
path = knowledge_service.get_entity_path(entity)
|
||||
assert path == Path(knowledge_service.base_path / "knowledge/concept/test-entity.md")
|
||||
assert path == Path(knowledge_service.base_path / "concept/test-entity.md")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -71,8 +72,8 @@ async def test_create_relations(knowledge_service: KnowledgeService, entity_serv
|
||||
# Create relation
|
||||
relations = [
|
||||
RelationSchema(
|
||||
from_id=entity1.id,
|
||||
to_id=entity2.id,
|
||||
from_id=entity1.path_id,
|
||||
to_id=entity2.path_id,
|
||||
relation_type="test_relation",
|
||||
context="Test context",
|
||||
)
|
||||
@@ -82,9 +83,9 @@ async def test_create_relations(knowledge_service: KnowledgeService, entity_serv
|
||||
assert len(updated_entities) == 2
|
||||
|
||||
# Verify files were updated
|
||||
for entity_id in [entity1.id, entity2.id]:
|
||||
entity = await entity_service.get_entity(entity_id)
|
||||
file_path = knowledge_service.get_entity_path(entity)
|
||||
for entity in [entity1, entity2]:
|
||||
found = await entity_service.get_by_path_id(entity.path_id)
|
||||
file_path = knowledge_service.get_entity_path(found)
|
||||
content, _ = await knowledge_service.file_service.read_file(file_path)
|
||||
assert "test_relation" in content
|
||||
|
||||
@@ -100,7 +101,7 @@ async def test_add_observations(knowledge_service: KnowledgeService):
|
||||
# Add observations
|
||||
observations = ["Test observation 1", "Test observation 2"]
|
||||
updated_entity = await knowledge_service.add_observations(
|
||||
entity.id, observations, "Test context"
|
||||
entity.path_id, observations, "Test context"
|
||||
)
|
||||
|
||||
# Verify observations in DB
|
||||
@@ -128,7 +129,7 @@ async def test_delete_entity(knowledge_service: KnowledgeService):
|
||||
assert await knowledge_service.file_service.exists(file_path)
|
||||
|
||||
# Delete entity
|
||||
success = await knowledge_service.delete_entity(entity.id)
|
||||
success = await knowledge_service.delete_entity(entity.path_id)
|
||||
assert success
|
||||
|
||||
# Verify file was deleted
|
||||
@@ -136,7 +137,7 @@ async def test_delete_entity(knowledge_service: KnowledgeService):
|
||||
|
||||
# Verify entity was deleted from DB
|
||||
with pytest.raises(EntityNotFoundError):
|
||||
await knowledge_service.entity_service.get_entity(entity.id)
|
||||
await knowledge_service.entity_service.get_by_path_id(entity.path_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -151,7 +152,7 @@ async def test_delete_multiple_entities(knowledge_service: KnowledgeService):
|
||||
entities.append(entity)
|
||||
|
||||
# Delete entities
|
||||
success = await knowledge_service.delete_entities([e.id for e in entities])
|
||||
success = await knowledge_service.delete_entities([e.path_id for e in entities])
|
||||
assert success
|
||||
|
||||
# Verify files were deleted
|
||||
@@ -159,7 +160,7 @@ async def test_delete_multiple_entities(knowledge_service: KnowledgeService):
|
||||
file_path = knowledge_service.get_entity_path(entity)
|
||||
assert not await knowledge_service.file_service.exists(file_path)
|
||||
with pytest.raises(EntityNotFoundError):
|
||||
await knowledge_service.entity_service.get_entity(entity.id)
|
||||
await knowledge_service.entity_service.get_by_path_id(entity.path_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -187,14 +188,14 @@ async def test_entity_not_found_error(knowledge_service: KnowledgeService):
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_on_creation_failure(knowledge_service: KnowledgeService, monkeypatch):
|
||||
"""Should clean up DB entity if file write fails."""
|
||||
entity_ids: List[int] = []
|
||||
entity_ids: List[str] = []
|
||||
|
||||
# Capture created entity ID
|
||||
original_create = knowledge_service.entity_service.create_entity
|
||||
|
||||
async def mock_create_entity(*args, **kwargs):
|
||||
entity = await original_create(*args, **kwargs)
|
||||
entity_ids.append(entity.id)
|
||||
entity_ids.append(entity.path_id)
|
||||
return entity
|
||||
|
||||
# Force file write to fail
|
||||
@@ -213,7 +214,7 @@ async def test_cleanup_on_creation_failure(knowledge_service: KnowledgeService,
|
||||
# Verify entity was cleaned up
|
||||
assert len(entity_ids) == 1
|
||||
with pytest.raises(EntityNotFoundError):
|
||||
await knowledge_service.entity_service.get_entity(entity_ids[0])
|
||||
await knowledge_service.entity_service.get_by_path_id(entity_ids[0])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -225,13 +226,8 @@ async def test_skip_failed_batch_operations(knowledge_service: KnowledgeService)
|
||||
EntitySchema(name="test-2", entity_type="test", description="Test entity 2"),
|
||||
]
|
||||
|
||||
created = await knowledge_service.create_entities(entities)
|
||||
assert len(created) == 2 # Middle one should fail but not stop processing
|
||||
|
||||
# Verify first and last were created
|
||||
names = [e.name for e in created]
|
||||
assert "test-1" in names
|
||||
assert "test-2" in names
|
||||
with pytest.raises(IntegrityError):
|
||||
await knowledge_service.create_entities(entities)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -248,8 +244,8 @@ async def test_update_relations_in_files(knowledge_service: KnowledgeService):
|
||||
# Create relation
|
||||
relations = [
|
||||
RelationSchema(
|
||||
from_id=entity1.id,
|
||||
to_id=entity2.id,
|
||||
from_id=entity1.path_id,
|
||||
to_id=entity2.path_id,
|
||||
relation_type="connects_to",
|
||||
context="Test connection",
|
||||
)
|
||||
|
||||
@@ -27,7 +27,7 @@ async def observation_service(observation_repository: ObservationRepository) ->
|
||||
async def test_entity(session_maker: async_sessionmaker[AsyncSession]) -> Entity:
|
||||
"""Create a test entity."""
|
||||
async with session_maker() as session:
|
||||
entity = Entity(entity_type="test", name="test", description="Test entity")
|
||||
entity = Entity(entity_type="test", name="test", description="Test entity", path_id="test/test")
|
||||
session.add(entity)
|
||||
await session.commit()
|
||||
return entity
|
||||
|
||||
@@ -4,9 +4,8 @@ import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from basic_memory.models import Entity as EntityModel
|
||||
from basic_memory.models import Entity, Relation
|
||||
from basic_memory.repository.relation_repository import RelationRepository
|
||||
from basic_memory.schemas import Entity, Relation
|
||||
from basic_memory.services.relation_service import RelationService
|
||||
|
||||
|
||||
@@ -27,17 +26,19 @@ async def relation_service(relation_repository: RelationRepository) -> RelationS
|
||||
@pytest_asyncio.fixture
|
||||
async def test_entities(
|
||||
session_maker: async_sessionmaker[AsyncSession],
|
||||
) -> tuple[EntityModel, EntityModel]:
|
||||
) -> tuple[Entity, Entity]:
|
||||
"""Create two test entities."""
|
||||
async with session_maker() as session:
|
||||
entity1 = EntityModel(
|
||||
entity1 = Entity(
|
||||
name="test_entity_1",
|
||||
entity_type="test",
|
||||
path_id="test/test_entity_1",
|
||||
description="Test entity 1",
|
||||
)
|
||||
entity2 = EntityModel(
|
||||
entity2 = Entity(
|
||||
name="test_entity_2",
|
||||
entity_type="test",
|
||||
path_id="test/test_entity_2",
|
||||
description="Test entity 2",
|
||||
)
|
||||
session.add_all([entity1, entity2])
|
||||
@@ -47,7 +48,7 @@ async def test_entities(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_relation(
|
||||
relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
|
||||
relation_service: RelationService, test_entities: tuple[Entity, Entity]
|
||||
):
|
||||
"""Test creating a basic relation between two entities."""
|
||||
entity1, entity2 = test_entities
|
||||
@@ -63,7 +64,7 @@ async def test_create_relation(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_relations(
|
||||
relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
|
||||
relation_service: RelationService, test_entities: tuple[Entity, Entity]
|
||||
):
|
||||
"""Test creating a basic relation between two entities."""
|
||||
entity1, entity2 = test_entities
|
||||
@@ -88,7 +89,7 @@ async def test_create_relations(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_relation_with_context(
|
||||
relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
|
||||
relation_service: RelationService, test_entities: tuple[Entity, Entity]
|
||||
):
|
||||
"""Test creating a relation with context information."""
|
||||
entity1, entity2 = test_entities
|
||||
@@ -104,7 +105,7 @@ async def test_create_relation_with_context(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_relation(
|
||||
relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
|
||||
relation_service: RelationService, test_entities: tuple[Entity, Entity]
|
||||
):
|
||||
"""Test deleting a relation between entities."""
|
||||
entity1, entity2 = test_entities
|
||||
@@ -113,73 +114,38 @@ async def test_delete_relation(
|
||||
relation_data = Relation(from_id=entity1.id, to_id=entity2.id, relation_type="test_relation")
|
||||
await relation_service.create_relation(relation_data)
|
||||
|
||||
# Create Entity schema instances for delete_relation call
|
||||
from_entity = Entity(
|
||||
id=entity1.id,
|
||||
name=entity1.name,
|
||||
entity_type=entity1.entity_type,
|
||||
description=entity1.description,
|
||||
observations=[],
|
||||
)
|
||||
to_entity = Entity(
|
||||
id=entity2.id,
|
||||
name=entity2.name,
|
||||
entity_type=entity2.entity_type,
|
||||
description=entity2.description,
|
||||
observations=[],
|
||||
)
|
||||
|
||||
# Delete the relation
|
||||
result = await relation_service.delete_relation(from_entity, to_entity, "test_relation")
|
||||
result = await relation_service.delete_relation(entity1, entity2, "test_relation")
|
||||
|
||||
assert result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_nonexistent_relation(
|
||||
relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
|
||||
relation_service: RelationService, test_entities: tuple[Entity, Entity]
|
||||
):
|
||||
"""Test trying to delete a relation that doesn't exist."""
|
||||
entity1, entity2 = test_entities
|
||||
|
||||
from_entity = Entity(
|
||||
id=entity1.id,
|
||||
name=entity1.name,
|
||||
entity_type=entity1.entity_type,
|
||||
description=entity1.description,
|
||||
observations=[],
|
||||
)
|
||||
to_entity = Entity(
|
||||
id=entity2.id,
|
||||
name=entity2.name,
|
||||
entity_type=entity2.entity_type,
|
||||
description=entity2.description,
|
||||
observations=[],
|
||||
)
|
||||
|
||||
result = await relation_service.delete_relation(from_entity, to_entity, "nonexistent_relation")
|
||||
result = await relation_service.delete_relation(entity1, entity2, "nonexistent_relation")
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_relations_by_criteria(
|
||||
relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
|
||||
relation_service: RelationService, test_entities: tuple[Entity, Entity]
|
||||
):
|
||||
"""Test deleting relations by criteria."""
|
||||
entity1, entity2 = test_entities
|
||||
|
||||
# Create test relations
|
||||
await relation_service.create_relation(
|
||||
Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation1")
|
||||
)
|
||||
await relation_service.create_relation(
|
||||
Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation2")
|
||||
)
|
||||
relation1 = Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation1")
|
||||
await relation_service.create_relation(relation1)
|
||||
relation2 = Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation2")
|
||||
await relation_service.create_relation(relation2)
|
||||
|
||||
# Delete relations matching criteria
|
||||
result = await relation_service.delete_relations(
|
||||
[{"from_id": entity1.id, "to_id": entity2.id, "relation_type": "relation1"}]
|
||||
)
|
||||
result = await relation_service.delete_relations([relation1, relation2])
|
||||
|
||||
assert result is True
|
||||
assert result == 2
|
||||
|
||||
Reference in New Issue
Block a user