diff --git a/src/basic_memory/repository.py b/src/basic_memory/repository.py index f13cdb77..15333e76 100644 --- a/src/basic_memory/repository.py +++ b/src/basic_memory/repository.py @@ -1,3 +1,4 @@ +"""Repository implementations for basic-memory models.""" from typing import Type, Optional, Any, Sequence from sqlalchemy import select, func, Select, Executable, inspect, Result, Column, and_ from sqlalchemy.exc import NoResultFound @@ -8,24 +9,7 @@ from basic_memory.models import Entity, Observation, Relation, Base class Repository[T: Base]: - """ - Generic repository pattern implementation for handling database operations. - Adapted for basic-memory with string IDs and memory-specific operations. - - Provides basic CRUD operations with an async session. - - :param session: Async database session from SQLAlchemy. - :param Model: Database model class. - - Example usage: - async with async_sessionmaker() as session: - entity_repo = Repository(session, Entity) - entity = await entity_repo.create({ - 'id': '20240102-some-entity', - 'name': 'Example Entity', - 'entity_type': 'concept' - }) - """ + """Base repository implementation with generic CRUD operations.""" def __init__(self, session: AsyncSession, Model: Type[T]): self.session = session @@ -33,37 +17,19 @@ class Repository[T: Base]: self.primary_key: Column[Any] = inspect(self.Model).mapper.primary_key[0] self.valid_columns = [column.key for column in inspect(self.Model).columns] - async def refresh(self, instance: T) -> None: - """ - Refresh the state of the given instance from the database. - - :param instance: Instance to refresh - """ - await self.session.refresh(instance) + async def refresh(self, instance: T, relationships: list[str] | None = None) -> None: + """Refresh instance and optionally specified relationships.""" + await self.session.refresh(instance, relationships or []) async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]: - """ - Fetches records from the database with pagination. - - :param skip: Number of records to skip. - :param limit: Maximum number of records to fetch. - :return: List containing the fetched records. - """ + """Fetch records from the database with pagination.""" result = await self.session.execute( select(self.Model).offset(skip).limit(limit) ) return result.scalars().all() async def find_by_id(self, entity_id: str) -> Optional[T]: - """ - Fetches an entity by its unique identifier asynchronously. - - :param entity_id: Unique identifier of the entity (string timestamp-based ID) - :return: The entity if found, otherwise None - - Example: - entity = await repository.find_by_id('20240102-example-entity') - """ + """Fetch an entity by its unique identifier.""" try: result = await self.session.execute( select(self.Model).filter(self.primary_key == entity_id) @@ -72,39 +38,22 @@ class Repository[T: Base]: except NoResultFound: return None - async def create(self, entity_data: dict) -> T: - """ - Creates a new entity in the database from the provided data dictionary. - - :param entity_data: A dictionary containing data to be inserted - :return: The created entity - - Example: - >>> entity_data = { - ... 'id': '20240102-example', - ... 'name': 'Example Entity', - ... 'entity_type': 'concept', - ... 'description': 'An example entity' - ... } - >>> new_entity = await repo.create(entity_data) + async def create(self, entity_data: dict, model: Type[Base] | None = None) -> T: + """Create a new entity in the database from the provided data. + + Args: + entity_data: Dictionary containing the data to insert + model: Optional model class to use (defaults to self.Model) """ + model = model or self.Model model_data = {k: v for k, v in entity_data.items() if k in self.valid_columns} - entity = self.Model(**model_data) + entity = model(**model_data) self.session.add(entity) await self.session.flush() return entity async def update(self, entity_id: str, entity_data: dict) -> Optional[T]: - """ - Updates an entity with given entity_id using the provided entity_data. - - :param entity_id: String ID of the entity to update - :param entity_data: Dictionary containing the data to update - :return: The updated entity or None if not found - - Example: - updated = await repository.update('20240102-example', {'description': 'Updated description'}) - """ + """Update an entity with the given data.""" try: result = await self.session.execute( select(self.Model).filter(self.primary_key == entity_id) @@ -119,15 +68,7 @@ class Repository[T: Base]: return None async def delete(self, entity_id: str) -> bool: - """ - Deletes an entity from the database. - - :param entity_id: String ID of the entity to delete - :return: Boolean indicating if the entity was deleted - - Example: - success = await repository.delete('20240102-example') - """ + """Delete an entity from the database.""" try: result = await self.session.execute( select(self.Model).filter(self.primary_key == entity_id) @@ -140,12 +81,7 @@ class Repository[T: Base]: return False async def count(self, query: Executable | None = None) -> int: - """ - Counts entities in the database table. - - :param query: Optional SQL query to modify the count operation - :return: Number of matching entities - """ + """Count entities in the database table.""" if query is None: query = select(func.count()).select_from(self.Model) result = await self.session.execute(query) @@ -153,130 +89,74 @@ class Repository[T: Base]: return scalar if scalar is not None else 0 async def execute_query(self, query: Executable) -> Result[Any]: - """ - Executes the given query asynchronously. - - :param query: An executable query instance - :return: Query result - """ + """Execute a query asynchronously.""" return await self.session.execute(query) async def find_one(self, query: Select[tuple[T]]) -> Optional[T]: - """ - Executes a query and retrieves a single record. - - :param query: The query to execute - :return: Single record or None - """ + """Execute a query and retrieve a single record.""" result = await self.execute_query(query) return result.scalars().one_or_none() class EntityRepository(Repository[Entity]): - """ - Repository for Entity model with memory-specific operations. - """ + """Repository for Entity model with memory-specific operations.""" async def find_by_id(self, entity_id: str) -> Optional[Entity]: - """ - Find entity by ID with all relationships eagerly loaded. - - Uses selectinload to eagerly load observations, outgoing and incoming relations - in a single query. This is necessary because: - 1. In async code, lazy loading relations after the session closes doesn't work - 2. Our service layer often needs the complete entity - 3. Using a single query with selectinload is more efficient - - :param entity_id: Entity ID to search for - :return: Entity if found with everything loaded, None otherwise - """ + """Find entity by ID with all relationships eagerly loaded.""" try: + # First load base entity result = await self.session.execute( - select(Entity) - .filter(Entity.id == entity_id) - .options( - selectinload(Entity.observations), - selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations) - ) + select(Entity).filter(Entity.id == entity_id) ) - return result.scalars().one() + entity = result.scalars().one() + + # Force refresh of all relationships + await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations']) + + return entity except NoResultFound: return None async def find_by_name(self, name: str) -> Optional[Entity]: - """ - Find an entity by its unique name. - - :param name: Entity name to search for - :return: Entity if found, None otherwise - """ + """Find an entity by its unique name.""" query = ( select(Entity) .filter(Entity.name == name) - .options( - selectinload(Entity.observations), - selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations) - ) ) - return await self.find_one(query) + result = await self.session.execute(query) + entity = result.scalars().one_or_none() + if entity: + await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations']) + return entity async def search_by_type(self, entity_type: str, skip: int = 0, limit: int = 100) -> Sequence[Entity]: - """ - Search for entities of a specific type. - - :param entity_type: Type to search for - :param skip: Number of records to skip - :param limit: Maximum records to return - :return: List of matching entities - """ + """Search for entities of a specific type.""" query = select(Entity).filter(Entity.entity_type == entity_type).offset(skip).limit(limit) result = await self.execute_query(query) return result.scalars().all() class ObservationRepository(Repository[Observation]): - """ - Repository for Observation model with memory-specific operations. - """ + """Repository for Observation model with memory-specific operations.""" async def find_by_entity(self, entity_id: str) -> Sequence[Observation]: - """ - Find all observations for a specific entity. - - :param entity_id: ID of the entity to find observations for - :return: List of observations - """ + """Find all observations for a specific entity.""" query = select(Observation).filter(Observation.entity_id == entity_id) result = await self.execute_query(query) return result.scalars().all() async def find_by_context(self, context: str) -> Sequence[Observation]: - """ - Find observations with a specific context. - - :param context: Context to search for - :return: List of matching observations - """ + """Find observations with a specific context.""" query = select(Observation).filter(Observation.context == context) result = await self.execute_query(query) return result.scalars().all() class RelationRepository(Repository[Relation]): - """ - Repository for Relation model with memory-specific operations. - """ + """Repository for Relation model with memory-specific operations.""" async def find_by_entities(self, from_id: str, to_id: str) -> Sequence[Relation]: - """ - Find all relations between two entities. - - :param from_id: Source entity ID - :param to_id: Target entity ID - :return: List of relations between the entities - """ + """Find all relations between two entities.""" query = select(Relation).filter( and_( Relation.from_id == from_id, @@ -287,12 +167,7 @@ class RelationRepository(Repository[Relation]): return result.scalars().all() async def find_by_type(self, relation_type: str) -> Sequence[Relation]: - """ - Find all relations of a specific type. - - :param relation_type: Type of relation to find - :return: List of matching relations - """ + """Find all relations of a specific type.""" query = select(Relation).filter(Relation.relation_type == relation_type) result = await self.execute_query(query) return result.scalars().all() \ No newline at end of file diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index 720cd5f6..4b0d2d6a 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -3,10 +3,12 @@ from datetime import datetime, UTC from pathlib import Path from basic_memory.repository import EntityRepository -from basic_memory.schemas import EntityIn -from basic_memory.models import Entity +from basic_memory.schemas import EntityIn, ObservationIn +from basic_memory.models import Entity, Observation +from basic_memory.fileio import EntityNotFoundError from . import ServiceError + class EntityService: """ Service for managing entities in the database. @@ -18,28 +20,35 @@ class EntityService: self.entity_repo = entity_repo async def create_entity(self, entity: EntityIn) -> Entity: - """Create a new entity in the database.""" - # Create DB record - db_data = { - **entity.model_dump(), + """Create a new entity in the database. + + Note: ID is generated by the EntityIn validator before reaching this method. + """ + # Create base entity first + base_data = { + "id": entity.id, # Include the generated ID + "name": entity.name, + "entity_type": entity.entity_type, "created_at": datetime.now(UTC), } - return await self.entity_repo.create(db_data) + created_entity = await self.entity_repo.create(base_data) + await self.entity_repo.refresh(created_entity, ['observations', 'outgoing_relations', 'incoming_relations']) + return created_entity async def get_entity(self, entity_id: str) -> Entity: """Get entity by ID.""" db_entity = await self.entity_repo.find_by_id(entity_id) if not db_entity: - raise ServiceError(f"Entity not found: {entity_id}") + raise EntityNotFoundError(f"Entity not found: {entity_id}") return db_entity - # TODO name is not uniaue + # TODO name is not unique async def get_by_name(self, name: str) -> Entity: """Get entity by name.""" db_entity = await self.entity_repo.find_by_name(name) if not db_entity: - raise ServiceError(f"Entity not found: {name}") + raise EntityNotFoundError(f"Entity not found: {name}") return db_entity diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index c400a534..56283616 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -29,9 +29,6 @@ class MemoryService: async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]: """Create multiple entities with their observations.""" entities_in = [EntityIn.model_validate(data) for data in entities_data] - print(f"\nCreating entities with observations:") - for e in entities_in: - print(f"Entity {e.name}: {len(e.observations)} observations") # Write files in parallel (filesystem is source of truth) async def write_file(entity: EntityIn): @@ -41,14 +38,11 @@ class MemoryService: await asyncio.gather(*file_writes) async def create_entity_in_db(entity_in: EntityIn): - print(f"\nCreating entity in DB: {entity_in.name}") - db_entity = await self.entity_service.create_entity(entity_in) - print(f"Adding {len(entity_in.observations)} observations to DB for {entity_in.name}") + await self.entity_service.create_entity(entity_in) await self.observation_service.add_observations(entity_in, entity_in.observations) [await self.relation_service.create_relation(relation_in) for relation_in in entity_in.relations] # query the entity again to return relations final_entity = await self.entity_service.get_entity(entity_in.id) - print(f"Final entity {final_entity.name} has {len(final_entity.observations)} observations in DB") return final_entity # Update database index sequentially @@ -91,27 +85,21 @@ class MemoryService: """ # Create new observations new_observations = ObservationsIn.model_validate(observations_in) - print(f"\nAdding new observations to entity {new_observations.entity_id}") - print(f"New observations to add: {len(new_observations.observations)}") # Read entity from filesystem entity = await read_entity_file(self.entities_path, new_observations.entity_id) - print(f"Entity {entity.id} from file has {len(entity.observations)} observations") - + # Create new observations for the entity for obs in new_observations.observations: entity.observations.append(obs) - print(f"After appending, entity has {len(entity.observations)} observations") # Write updated entity file await write_entity_file(self.entities_path, entity) # Update database index added_observations = await self.observation_service.add_observations(entity, new_observations.observations) - print(f"Added {len(added_observations)} observations to DB") - + db_entity = await self.entity_service.get_entity(entity.id) - print(f"Entity {entity.id} in DB now has {len(db_entity.observations)} observations") return added_observations async def delete_entities(self, entity_names: List[str]) -> None: diff --git a/src/basic_memory/services/observation_service.py b/src/basic_memory/services/observation_service.py index 1ed82ae6..0ced1cd8 100644 --- a/src/basic_memory/services/observation_service.py +++ b/src/basic_memory/services/observation_service.py @@ -25,23 +25,31 @@ class ObservationService: Add multiple observations to an entity. Returns the created observations with IDs set. """ - print(f"\nObservationService.add_observations called for entity {entity.id}") - print(f"Adding {len(observations)} observations") - async def add_observation(observation: ObservationIn) -> Observation: try: - return await self.observation_repo.create({ + obs = await self.observation_repo.create({ 'entity_id': entity.id, 'content': observation.content, 'context': observation.context, 'created_at': datetime.now(UTC) }) + # Ensure each observation is flushed + await self.observation_repo.session.flush() + # Refresh to get latest state + await self.observation_repo.session.refresh(obs) + return obs except Exception as e: raise DatabaseSyncError(f"Failed to add observation to database: {str(e)}") from e # Add each observation and collect the results created_observations = [await add_observation(obs) for obs in observations] - print(f"Created {len(created_observations)} observations in DB") + + # Make sure observations are in sync before returning + # This helps ensure related entities see the new observations + await self.observation_repo.session.flush() + for obs in created_observations: + await self.observation_repo.session.refresh(obs) + return created_observations async def search_observations(self, query: str) -> List[Observation]: diff --git a/src/basic_memory/services/relation_service.py b/src/basic_memory/services/relation_service.py index 72a59f34..7162be16 100644 --- a/src/basic_memory/services/relation_service.py +++ b/src/basic_memory/services/relation_service.py @@ -25,8 +25,7 @@ class RelationService: try: db_data = relation.model_dump() db_data['created_at'] = datetime.now(UTC) - await self.relation_repo.create(db_data) - return relation + return await self.relation_repo.create(db_data) except Exception as e: raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e diff --git a/tests/conftest.py b/tests/conftest.py index f511ab10..cbcbb7fd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,4 @@ """Common test fixtures for basic-memory.""" -import pytest import pytest_asyncio from pathlib import Path import tempfile @@ -16,6 +15,8 @@ from basic_memory.deps import ( get_relation_service, get_memory_service ) +from basic_memory.schemas import EntityIn + @pytest_asyncio.fixture(scope="function") async def engine(): @@ -38,7 +39,7 @@ async def engine(): @pytest_asyncio.fixture(scope="function") async def session(engine): """Create an async session factory and yield a session""" - async_session = async_sessionmaker(engine, expire_on_commit=False) + async_session = async_sessionmaker(engine) # Removed expire_on_commit=False async with async_session() as session: try: yield session @@ -104,7 +105,8 @@ async def memory_service( @pytest_asyncio.fixture async def test_entity(entity_service): """Create a test entity for reuse in tests.""" - return await entity_service.create_entity( + entity_data = EntityIn( name="Test Entity", entity_type="test", - ) \ No newline at end of file + ) + return await entity_service.create_entity(entity_data) \ No newline at end of file diff --git a/tests/test_entity_service.py b/tests/test_entity_service.py index cb41c46e..30feebee 100644 --- a/tests/test_entity_service.py +++ b/tests/test_entity_service.py @@ -1,105 +1,44 @@ +"""Tests for EntityService.""" import pytest import pytest_asyncio -from datetime import datetime from pathlib import Path import tempfile from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker from sqlalchemy.pool import StaticPool from basic_memory.fileio import EntityNotFoundError -from basic_memory.models import Entity as DbEntity, Base +from basic_memory.models import Entity as DbEntity, Base, Entity from basic_memory.repository import EntityRepository -from basic_memory.services import EntityService, ServiceError, DatabaseSyncError -from basic_memory.schemas import Entity, Observation +from basic_memory.schemas import EntityIn, ObservationIn +from basic_memory.services import EntityService pytestmark = pytest.mark.asyncio -@pytest_asyncio.fixture(scope="function") -async def engine(): - """Create an async engine using in-memory SQLite database""" - engine = create_async_engine( - "sqlite+aiosqlite:///:memory:", # In-memory database - echo=False, # Set to True for SQL logging - poolclass=StaticPool, # Ensure single connection for in-memory db - connect_args={"check_same_thread": False} # Allow multi-threaded access - ) - - # Create all tables - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) - - try: - yield engine - finally: - await engine.dispose() - -@pytest_asyncio.fixture(scope="function") -async def session(engine): - """Create an async session factory and yield a session""" - async_session = async_sessionmaker(engine, expire_on_commit=False) - async with async_session() as session: - try: - yield session - await session.commit() - except Exception: - await session.rollback() - raise - -@pytest_asyncio.fixture -async def entity_repo(session): - """Create an EntityRepository instance.""" - return EntityRepository(session, DbEntity) - -@pytest_asyncio.fixture -async def entity_service(session, entity_repo): - """Fixture providing initialized EntityService with temp directories.""" - with tempfile.TemporaryDirectory() as temp_dir: - project_path = Path(temp_dir) / "test-project" - entities_path = project_path / "entities" - entities_path.mkdir(parents=True) - - service = EntityService(project_path, entity_repo) - yield service - -# Happy Path Tests async def test_create_entity_success(entity_service): """Test successful entity creation.""" - observations = ["First observation", "Second observation"] - - # Act - entity = await entity_service.create_entity( + entity_data = EntityIn( name="Test Entity", entity_type="test", - observations=observations ) + # Act + entity = await entity_service.create_entity(entity_data) + # Assert Entity assert isinstance(entity, Entity) assert entity.name == "Test Entity" assert entity.entity_type == "test" - assert len(entity.observations) == 2 - assert entity.observations[0].content == "First observation" - assert entity.observations[1].content == "Second observation" - - # Verify file was created - entity_file = entity_service.entities_path / f"{entity.id}.md" - assert entity_file.exists() - content = entity_file.read_text() - assert "# Test Entity" in content - assert "type: test" in content - assert "First observation" in content - assert "Second observation" in content + assert entity.created_at is not None async def test_get_entity_success(entity_service): """Test successful entity retrieval.""" # Arrange - observations = ["Test observation"] - created = await entity_service.create_entity( + entity_data = EntityIn( name="Test Entity", entity_type="test", - observations=observations ) + created = await entity_service.create_entity(entity_data) # Act retrieved = await entity_service.get_entity(created.id) @@ -109,26 +48,22 @@ async def test_get_entity_success(entity_service): assert retrieved.id == created.id assert retrieved.name == created.name assert retrieved.entity_type == created.entity_type - assert len(retrieved.observations) == 1 - assert retrieved.observations[0].content == "Test observation" + # relations are tested in test_memory_service async def test_delete_entity_success(entity_service): """Test successful entity deletion.""" # Arrange - entity = await entity_service.create_entity( + entity_data = EntityIn( name="Test Entity", entity_type="test", - observations=["Test observation"] ) - entity_file = entity_service.entities_path / f"{entity.id}.md" - assert entity_file.exists() - + entity = await entity_service.create_entity(entity_data) + # Act result = await entity_service.delete_entity(entity.id) # Assert assert result is True - assert not entity_file.exists() with pytest.raises(EntityNotFoundError): await entity_service.get_entity(entity.id) @@ -146,13 +81,14 @@ async def test_create_entity_db_error(entity_service, monkeypatch): raise Exception("Mock DB error") monkeypatch.setattr(entity_service.entity_repo, "create", mock_create) - # Act/Assert - both file and DB operations should fail + entity_data = EntityIn( + name="Test Entity", + entity_type="test", + ) + + # Act/Assert with pytest.raises(Exception, match="Mock DB error"): - await entity_service.create_entity( - name="Test Entity", - entity_type="test", - observations=["Test observation"] - ) + await entity_service.create_entity(entity_data) async def test_delete_nonexistent_entity(entity_service): """Test deleting an entity that doesn't exist.""" @@ -164,94 +100,24 @@ async def test_delete_nonexistent_entity(entity_service): async def test_create_entity_with_special_chars(entity_service): """Test entity creation with special characters in name.""" name = "Test & Entity! With @ Special #Chars" - entity = await entity_service.create_entity( - name=name, + entity_data = EntityIn( + name=name, entity_type="test", - observations=["Test observation"] ) + entity = await entity_service.create_entity(entity_data) assert entity.name == name - entity_file = entity_service.entities_path / f"{entity.id}.md" - assert entity_file.exists() -async def test_create_entity_atomic_file_write(entity_service): - """Test that file writing is atomic (uses temp file).""" - # Act - entity = await entity_service.create_entity( - name="Test Entity", - entity_type="test", - observations=["Test observation"] - ) - - # Assert - temp_file = entity_service.entities_path / f"{entity.id}.md.tmp" - assert not temp_file.exists() # Temp file should be cleaned up - - entity_file = entity_service.entities_path / f"{entity.id}.md" - assert entity_file.exists() # Final file should exist - -async def test_rebuild_index(entity_service): - """Test rebuilding database index from filesystem.""" - # Arrange - Create some entities - entity1 = await entity_service.create_entity( - name="Test 1", - entity_type="test", - observations=["Test observation 1"] - ) - entity2 = await entity_service.create_entity( - name="Test 2", - entity_type="test", - observations=["Test observation 2"] - ) - - # Act - Rebuild index - await entity_service.rebuild_index() - - # Assert - Entities should be retrievable - retrieved1 = await entity_service.get_entity(entity1.id) - retrieved2 = await entity_service.get_entity(entity2.id) - - assert isinstance(retrieved1, Entity) - assert isinstance(retrieved2, Entity) - assert retrieved1.name == entity1.name - assert retrieved2.name == entity2.name - assert retrieved1.observations[0].content == "Test observation 1" - assert retrieved2.observations[0].content == "Test observation 2" - -# Tests for Pydantic Model Behavior async def test_entity_id_generation(entity_service): """Test that entities get unique IDs generated correctly.""" - entity = await entity_service.create_entity( + entity_data = EntityIn( name="Test Entity", - entity_type="test" + entity_type="test", + observations=[] ) + entity = await entity_service.create_entity(entity_data) + assert entity.id # ID should be generated - assert "-test-entity-" in entity.id # Should contain normalized name - assert len(entity.id.split("-")[-1]) == 8 # UUID part should be 8 chars - -async def test_entity_with_no_observations(entity_service): - """Test entity creation with no observations.""" - entity = await entity_service.create_entity( - name="Test Entity", - entity_type="test" - ) - - assert isinstance(entity, Entity) - assert entity.observations == [] - - # Check file format - entity_file = entity_service.entities_path / f"{entity.id}.md" - content = entity_file.read_text() - assert "## Observations" in content - assert content.strip().endswith("## Observations") # No observations after header - -# TODO: Add tests for: -# - Concurrent operations (using asyncio.gather) -# - Filesystem permissions issues -# - System crash simulation -# - Markdown formatting edge cases -# - Very long entity names/content -# - Unicode/special character handling -# - File system space issues \ No newline at end of file + assert "-test-entity" in entity.id # Should contain normalized name diff --git a/tests/test_observation_service.py b/tests/test_observation_service.py index c9ab5d6b..907d3091 100644 --- a/tests/test_observation_service.py +++ b/tests/test_observation_service.py @@ -1,38 +1,34 @@ """Tests for ObservationService.""" import pytest -import pytest_asyncio -from sqlalchemy import delete -from basic_memory.models import Observation as DbObservation -from basic_memory.repository import ObservationRepository -from basic_memory.services import ( - EntityService, ObservationService, - ServiceError, DatabaseSyncError -) -from basic_memory.schemas import Entity, Observation -from basic_memory.fileio import read_entity_file, FileOperationError +from basic_memory.models import Observation +from basic_memory.schemas import EntityIn, ObservationIn pytestmark = pytest.mark.asyncio async def test_add_observation_success(observation_service, test_entity): """Test successful observation addition.""" - # Act - observation = await observation_service.add_observation( - entity=test_entity, + observation_data = ObservationIn( content="New observation", context="test-context" ) + entity_data = EntityIn( + name=test_entity.name, + entity_type=test_entity.entity_type, + id=test_entity.id, + observations=[] + ) + + # Act + observations = await observation_service.add_observations(entity_data, [observation_data]) + # Assert - assert isinstance(observation, Observation) - assert observation.content == "New observation" - - # Verify file update - entity = await read_entity_file(observation_service.entities_path, test_entity.id) - assert len(entity.observations) == 1 - assert any(obs.content == "New observation" for obs in entity.observations) - + assert len(observations) == 1 + assert isinstance(observations[0], Observation) + assert observations[0].content == "New observation" + # Verify database index db_observations = await observation_service.observation_repo.find_by_entity(test_entity.id) assert len(db_observations) == 1 @@ -40,27 +36,22 @@ async def test_add_observation_success(observation_service, test_entity): for obs in db_observations) -async def test_file_operation_error(observation_service, test_entity, mocker): - """Test handling of file operation errors.""" - async def mock_write(*args, **kwargs): - print("Mock write called with:", args, kwargs) - raise FileOperationError("Mock file error") - - mocker.patch('basic_memory.services.observation_service.write_entity_file', mock_write) - - with pytest.raises(FileOperationError): - await observation_service.add_observation( - test_entity, - "Test observation" - ) - async def test_search_observations(observation_service, test_entity): """Test searching observations across entities.""" # Arrange - await observation_service.add_observation(test_entity, "Unique test content") - await observation_service.add_observation(test_entity, "Other content") + entity_data = EntityIn( + name=test_entity.name, + entity_type=test_entity.entity_type, + id=test_entity.id, + observations=[] + ) + await observation_service.add_observations( + entity_data, + [ObservationIn(content="Unique test content"), ObservationIn(content="Other content")] + ) + # Act results = await observation_service.search_observations("unique") @@ -72,17 +63,19 @@ async def test_search_observations(observation_service, test_entity): async def test_get_observations_by_context(observation_service, test_entity): """Test retrieving observations by context.""" # Arrange - await observation_service.add_observation( - test_entity, - "Context observation", - context="test-context" - ) - await observation_service.add_observation( - test_entity, - "Other observation", - context="other-context" + entity_data = EntityIn( + name=test_entity.name, + entity_type=test_entity.entity_type, + id=test_entity.id, + observations=[] ) + await observation_service.add_observations( + entity_data, + [ObservationIn(content="Context observation", context="test-context"), + ObservationIn(content="Other observation", context="other-context")] + ) + # Act results = await observation_service.get_observations_by_context("test-context") @@ -91,53 +84,31 @@ async def test_get_observations_by_context(observation_service, test_entity): assert results[0].content == "Context observation" -async def test_rebuild_observation_index(observation_service, test_entity): - """Test rebuilding observation index from filesystem.""" - # Arrange - Add observations and clear database - await observation_service.add_observation(test_entity, "Test observation 1") - await observation_service.add_observation(test_entity, "Test observation 2") - - # Clear database but keep files - await observation_service.observation_repo.execute_query(delete(DbObservation)) - - # Act - await observation_service.rebuild_observation_index() - - # Assert - db_observations = await observation_service.observation_repo.find_by_entity(test_entity.id) - assert len(db_observations) == 2 - observation_contents = {obs.content for obs in db_observations} - assert observation_contents == { - "Test observation 1", - "Test observation 2" - } - - # Edge Cases async def test_observation_with_special_characters(observation_service, test_entity): """Test handling observations with special characters.""" content = "Test & observation with @#$% special chars!" - observation = await observation_service.add_observation( - test_entity, - content + entity_data = EntityIn( + name=test_entity.name, + entity_type=test_entity.entity_type, + id=test_entity.id, ) - assert observation.content == content - - # Verify file content - entity = await read_entity_file(observation_service.entities_path, test_entity.id) - assert any(obs.content == content for obs in entity.observations) + + observations = await observation_service.add_observations(entity_data, [ObservationIn(content=content)]) + assert observations[0].content == content async def test_very_long_observation(observation_service, test_entity): """Test handling very long observation content.""" long_content = "Very long observation " * 100 # ~1800 characters - observation = await observation_service.add_observation( - test_entity, - long_content + entity_data = EntityIn( + name=test_entity.name, + entity_type=test_entity.entity_type, + id=test_entity.id, + observations=[] ) - assert observation.content == long_content - - # Verify file content - entity = await read_entity_file(observation_service.entities_path, test_entity.id) - assert any(obs.content.rstrip() == long_content.rstrip() for obs in entity.observations) \ No newline at end of file + + observations = await observation_service.add_observations(entity_data, [ObservationIn(content=long_content)]) + assert observations[0].content == long_content + \ No newline at end of file diff --git a/tests/test_relation_service.py b/tests/test_relation_service.py index 77f44863..b5efa97d 100644 --- a/tests/test_relation_service.py +++ b/tests/test_relation_service.py @@ -1,16 +1,9 @@ """Tests for RelationService.""" import pytest import pytest_asyncio -from sqlalchemy import delete, select -from basic_memory.models import Relation as DbRelation -from basic_memory.repository import EntityRepository, RelationRepository -from basic_memory.services import ( - EntityService, RelationService, - ServiceError, DatabaseSyncError, RelationError -) -from basic_memory.schemas import Entity, Relation -from basic_memory.fileio import read_entity_file, FileOperationError, write_entity_file +from basic_memory.fileio import FileOperationError +from basic_memory.schemas import EntityIn, RelationIn pytestmark = pytest.mark.asyncio @@ -18,14 +11,20 @@ pytestmark = pytest.mark.asyncio @pytest_asyncio.fixture async def sample_entities(entity_service): """Create two sample entities for testing relations""" - entity1 = await entity_service.create_entity( + entity1_data = EntityIn( name="test_entity_1", - entity_type="test_type" + entity_type="test_type", + observations=[], + relations=[] ) - entity2 = await entity_service.create_entity( + entity2_data = EntityIn( name="test_entity_2", - entity_type="test_type" + entity_type="test_type", + observations=[], + relations=[] ) + entity1 = await entity_service.create_entity(entity1_data) + entity2 = await entity_service.create_entity(entity2_data) return entity1, entity2 @@ -38,29 +37,19 @@ async def test_create_relation(relation_service, sample_entities): """Test creating a basic relation between two entities""" entity1, entity2 = sample_entities - relation = await relation_service.create_relation( - from_entity=entity1, - to_entity=entity2, + relation_data = RelationIn( + from_id=entity1.id, + to_id=entity2.id, relation_type="test_relation" ) - # Check Entity objects in relation - assert relation.from_entity.id == entity1.id - assert relation.to_entity.id == entity2.id + relation = await relation_service.create_relation(relation_data) + + # Check relation was created correctly + assert relation.from_id == entity1.id + assert relation.to_id == entity2.id assert relation.relation_type == "test_relation" - - # Verify relation was added to source entity's relations - assert hasattr(entity1, 'relations') - assert len(entity1.relations) == 1 - assert entity1.relations[0].id == relation.id - - # Verify file was written with relation - entity_file = relation_service.entities_path / f"{entity1.id}.md" - assert entity_file.exists() - content = entity_file.read_text() - assert "## Relations" in content - assert f"[{entity2.id}] test_relation" in content - + # Verify database was updated with correct IDs db_relation = await relation_service.relation_repo.find_by_id(relation.id) assert db_relation is not None @@ -69,44 +58,22 @@ async def test_create_relation(relation_service, sample_entities): assert db_relation.relation_type == "test_relation" -async def test_file_operation_error(relation_service, sample_entities, mocker): - """Test handling of file operation errors.""" - entity1, entity2 = sample_entities - - # Add debug to see if mock is being called - async def mock_write(*args, **kwargs): - print("Mock write called with:", args, kwargs) - raise FileOperationError("Mock file error") - - # Patch where the function is used, not where it's imported from - mocker.patch('basic_memory.services.relation_service.write_entity_file', mock_write) - - with pytest.raises(FileOperationError): - await relation_service.create_relation( - from_entity=entity1, - to_entity=entity2, - relation_type="test_relation" - ) - async def test_create_relation_with_context(relation_service, sample_entities): """Test creating a relation with context information""" entity1, entity2 = sample_entities - relation = await relation_service.create_relation( - from_entity=entity1, - to_entity=entity2, + relation_data = RelationIn( + from_id=entity1.id, + to_id=entity2.id, relation_type="test_relation", context="test context" ) + relation = await relation_service.create_relation(relation_data) + assert relation.context == "test context" - - # Verify context in file - entity_file = relation_service.entities_path / f"{entity1.id}.md" - content = entity_file.read_text() - assert f"[{entity2.id}] test_relation | test context" in content - + # Verify context in database db_relation = await relation_service.relation_repo.find_by_id(relation.id) assert db_relation.context == "test context" \ No newline at end of file diff --git a/tests/test_repository.py b/tests/test_repository.py index 65b72c1e..0154458c 100644 --- a/tests/test_repository.py +++ b/tests/test_repository.py @@ -123,7 +123,6 @@ class TestObservationRepository: async def sample_observation(self, observation_repository: ObservationRepository, sample_entity: Entity): """Create a sample observation for testing""" observation_data = { - 'id': '20240102-test-obs', 'entity_id': sample_entity.id, 'content': 'Test observation', 'context': 'test-context' @@ -137,16 +136,15 @@ class TestObservationRepository: ): """Test creating a new observation""" observation_data = { - 'id': '20240102-obs', 'entity_id': sample_entity.id, 'content': 'Test content', 'context': 'test-context' } observation = await observation_repository.create(observation_data) - assert observation.id == '20240102-obs' assert observation.entity_id == sample_entity.id assert observation.content == 'Test content' + assert observation.id is not None # Should be auto-generated async def test_find_by_entity( self, @@ -193,7 +191,6 @@ class TestRelationRepository: ): """Create a sample relation for testing""" relation_data = { - 'id': '20240102-test-rel', 'from_id': sample_entity.id, 'to_id': related_entity.id, 'relation_type': 'test_relation', @@ -209,7 +206,6 @@ class TestRelationRepository: ): """Test creating a new relation""" relation_data = { - 'id': '20240102-rel', 'from_id': sample_entity.id, 'to_id': related_entity.id, 'relation_type': 'test_relation', @@ -217,10 +213,10 @@ class TestRelationRepository: } relation = await relation_repository.create(relation_data) - assert relation.id == '20240102-rel' assert relation.from_id == sample_entity.id assert relation.to_id == related_entity.id assert relation.relation_type == 'test_relation' + assert relation.id is not None # Should be auto-generated async def test_find_by_entities( self,