From 955fe899ab0ab02094661c00c9ad329eed49ea17 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 22 Dec 2024 12:17:59 -0600 Subject: [PATCH] fix repository tests --- src/basic_memory/models/knowledge.py | 1 + .../repository/observation_repository.py | 4 +--- .../repository/relation_repository.py | 4 ++-- tests/repository/test_observation_repository.py | 15 ++++++--------- tests/repository/test_relation_repository.py | 5 ----- 5 files changed, 10 insertions(+), 19 deletions(-) diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index e33ca7ae..76dd4f43 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -87,6 +87,7 @@ class Observation(Base): id: Mapped[int] = mapped_column(Integer, primary_key=True) entity_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id")) content: Mapped[str] = mapped_column(Text) + context: Mapped[str] = mapped_column(Text, nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) # Relationships diff --git a/src/basic_memory/repository/observation_repository.py b/src/basic_memory/repository/observation_repository.py index a62a50dd..bd502eee 100644 --- a/src/basic_memory/repository/observation_repository.py +++ b/src/basic_memory/repository/observation_repository.py @@ -2,7 +2,6 @@ from typing import Sequence -from loguru import logger from sqlalchemy import select from sqlalchemy.ext.asyncio import async_sessionmaker @@ -15,9 +14,8 @@ class ObservationRepository(Repository[Observation]): def __init__(self, session_maker: async_sessionmaker): super().__init__(session_maker, Observation) - logger.debug("Initialized ObservationRepository") - async def find_by_entity(self, entity_id: str) -> Sequence[Observation]: + async def find_by_entity(self, entity_id: int) -> Sequence[Observation]: """Find all observations for a specific entity.""" query = select(Observation).filter(Observation.entity_id == entity_id) result = await self.execute_query(query) diff --git a/src/basic_memory/repository/relation_repository.py b/src/basic_memory/repository/relation_repository.py index f7c2c3c9..afa41d8b 100644 --- a/src/basic_memory/repository/relation_repository.py +++ b/src/basic_memory/repository/relation_repository.py @@ -15,13 +15,13 @@ class RelationRepository(Repository[Relation]): def __init__(self, session_maker: async_sessionmaker): super().__init__(session_maker, Relation) - async def find_by_entity(self, from_entity_id: str) -> Sequence[Relation]: + 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) result = await self.execute_query(query) return result.scalars().all() - async def find_by_entities(self, from_id: str, to_id: str) -> Sequence[Relation]: + async def find_by_entities(self, from_id: int, to_id: int) -> Sequence[Relation]: """Find all relations between two entities.""" query = select(Relation).filter(and_(Relation.from_id == from_id, Relation.to_id == to_id)) result = await self.execute_query(query) diff --git a/tests/repository/test_observation_repository.py b/tests/repository/test_observation_repository.py index 1c4ff99b..80f45019 100644 --- a/tests/repository/test_observation_repository.py +++ b/tests/repository/test_observation_repository.py @@ -87,10 +87,9 @@ 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( - id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" - ) + entity = Entity(name="test_entity", entity_type="test", description="Test entity") session.add(entity) + await session.flush() # Create test observations obs1 = Observation(entity_id=entity.id, content="Test observation 1") @@ -111,10 +110,9 @@ 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( - id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" - ) + entity = Entity(name="test_entity", entity_type="test", description="Test entity") session.add(entity) + await session.flush() # Create test observation obs = Observation(entity_id=entity.id, content="Test observation") @@ -134,10 +132,9 @@ 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( - id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" - ) + entity = Entity(name="test_entity", entity_type="test", description="Test entity") session.add(entity) + await session.flush() # Create test observations obs1 = Observation(entity_id=entity.id, content="Delete this observation") diff --git a/tests/repository/test_relation_repository.py b/tests/repository/test_relation_repository.py index 455631cc..c56793a2 100644 --- a/tests/repository/test_relation_repository.py +++ b/tests/repository/test_relation_repository.py @@ -13,7 +13,6 @@ from basic_memory.repository.relation_repository import RelationRepository async def source_entity(session_maker): """Create a source entity for testing relations.""" entity = Entity( - id="source/test_entity", name="test_source", entity_type="source", description="Source entity", @@ -28,7 +27,6 @@ async def source_entity(session_maker): async def target_entity(session_maker): """Create a target entity for testing relations.""" entity = Entity( - id="target/test_entity", name="test_target", entity_type="target", description="Target entity", @@ -56,7 +54,6 @@ async def test_relations(session_maker, source_entity, target_entity): async def related_entity(entity_repository): """Create a second entity for testing relations""" entity_data = { - "id": "20240102-related", "name": "Related Entity", "entity_type": "test", "description": "A related test entity", @@ -229,14 +226,12 @@ async def test_delete_by_fields_all_fields( from_id=relation.from_id, # pyright: ignore [reportArgumentType] to_id=relation.to_id, # pyright: ignore [reportArgumentType] relation_type=relation.relation_type, # pyright: ignore [reportArgumentType] - context=relation.context, # pyright: ignore [reportArgumentType] ) assert result is True # Verify only exact match was deleted remaining = await relation_repository.find_by_type(relation.relation_type) assert len(remaining) == 1 # One other relation_one should remain - assert remaining[0].context != relation.context @pytest.mark.asyncio