fix tests

This commit is contained in:
phernandez
2024-12-24 20:29:21 -06:00
parent d713097333
commit 90ca41fead
16 changed files with 137 additions and 138 deletions
@@ -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}
)