diff --git a/src/basic_memory/db.py b/src/basic_memory/db.py index b1b4cadb..f4bb8abd 100644 --- a/src/basic_memory/db.py +++ b/src/basic_memory/db.py @@ -1,90 +1,101 @@ """Database configuration and initialization for basic-memory.""" + +from contextlib import asynccontextmanager from enum import Enum from pathlib import Path from typing import AsyncGenerator -from contextlib import asynccontextmanager from loguru import logger -from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncEngine, AsyncSession +from sqlalchemy import text +from sqlalchemy.ext.asyncio import ( + create_async_engine, + async_sessionmaker, + AsyncEngine, + AsyncSession, +) from sqlalchemy.pool import StaticPool from basic_memory.models import Base + class DatabaseType(Enum): """Types of database configurations.""" - MEMORY = "memory" # In-memory SQLite for testing - FILESYSTEM = "file" # File-based SQLite for projects -def get_database_url(project_path: Path, db_type: DatabaseType, ) -> str: + MEMORY = "memory" # In-memory SQLite for testing + FILESYSTEM = "file" # File-based SQLite for projects + + +def get_database_url( + project_path: Path, + db_type: DatabaseType, +) -> str: """ Get database URL based on type and optional project path. - + Args: db_type: Type of database to configure project_path: Project directory for file-based DBs (required if type is FILESYSTEM) - + Returns: Database URL string - + Raises: ValueError: If project_path is required but not provided """ match db_type: case DatabaseType.MEMORY: return "sqlite+aiosqlite:///:memory:" - + case DatabaseType.FILESYSTEM: if not project_path: raise ValueError("project_path required for filesystem database") - + # Ensure data directory exists data_dir = project_path / "data" data_dir.mkdir(parents=True, exist_ok=True) - + db_path = data_dir / "memory.db" return f"sqlite+aiosqlite:///{db_path}" + async def init_database(url: str, echo: bool = False) -> tuple[AsyncEngine, async_sessionmaker]: """ Initialize database with schema. - + Args: url: Database URL echo: Whether to echo SQL statements - + Returns: Configured async engine and session factory """ # Configure engine based on URL connect_args = {"check_same_thread": False} + if url == "sqlite+aiosqlite:///:memory:": engine = create_async_engine( url, echo=echo, poolclass=StaticPool, # Single connection for in-memory - connect_args=connect_args + connect_args=connect_args, ) else: - engine = create_async_engine( - url, - echo=echo, - connect_args=connect_args - ) + engine = create_async_engine(url, echo=echo, connect_args=connect_args) # Create tables async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) # Create session factory for this engine - session_factory = async_sessionmaker( - engine, - expire_on_commit=False - ) - + session_factory = async_sessionmaker(engine, expire_on_commit=False) + return engine, session_factory + @asynccontextmanager -async def engine_session_factory(project_path: Path, db_type=DatabaseType.FILESYSTEM) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: +async def engine_session_factory( + project_path: Path, db_type=DatabaseType.FILESYSTEM +) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: """Get database engine and session factory with proper lifecycle management.""" url = get_database_url(project_path, db_type=db_type) engine, session_factory = await init_database(url) @@ -94,20 +105,25 @@ async def engine_session_factory(project_path: Path, db_type=DatabaseType.FILESY finally: await engine.dispose() + @asynccontextmanager -async def session(session_factory: async_sessionmaker[AsyncSession]) -> AsyncGenerator[AsyncSession, None]: +async def session( + session_factory: async_sessionmaker[AsyncSession], +) -> AsyncGenerator[AsyncSession, None]: """ Get database session with proper lifecycle management. - + Args: session_factory: Async session factory to create session from - + Yields: AsyncSession configured for engine """ # Create and yield session session = session_factory() try: + # Ensure foreign keys enabled for this session + await session.execute(text("PRAGMA foreign_keys=ON")) yield session await session.commit() except Exception: @@ -116,11 +132,12 @@ async def session(session_factory: async_sessionmaker[AsyncSession]) -> AsyncGen finally: await session.close() + async def dispose_database(engine: AsyncEngine): """ Clean up database engine. - + Args: engine: Engine to dispose """ - await engine.dispose() \ No newline at end of file + await engine.dispose() diff --git a/src/basic_memory/models.py b/src/basic_memory/models.py index 4af19b76..24ad7f35 100644 --- a/src/basic_memory/models.py +++ b/src/basic_memory/models.py @@ -5,6 +5,8 @@ from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, Unique from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase from sqlalchemy.ext.asyncio import AsyncAttrs +from basic_memory.utils import normalize_entity_id + class Base(AsyncAttrs, DeclarativeBase): """Base class for all models""" @@ -68,8 +70,8 @@ class Entity(Base): @classmethod def generate_id(cls, entity_type: str, name: str) -> str: """Generate a filesystem path-based ID for this entity.""" - # Normalize name for filesystem (handle spaces, special chars etc) - safe_name = name.lower().replace(" ", "_") + # Use common normalization for filesystem safety + safe_name = normalize_entity_id(name) return f"{entity_type}/{safe_name}" def get_file_path(self) -> str: diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 21974b79..babad5da 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -3,11 +3,16 @@ from typing import List, Optional, Annotated, TypeAlias from annotated_types import Len -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, BeforeValidator + +from basic_memory.utils import normalize_entity_id # Base Models Observation: TypeAlias = str +# Custom field types with validation +EntityId = Annotated[str, BeforeValidator(normalize_entity_id)] + class Relation(BaseModel): """ @@ -15,8 +20,8 @@ class Relation(BaseModel): Relations are always stored in active voice (e.g. "created", "teaches", etc.) """ - from_id: str - to_id: str + from_id: EntityId + to_id: EntityId relation_type: str context: Optional[str] = None @@ -28,7 +33,7 @@ class Entity(BaseModel): associated observations. """ - id: Optional[str] = None + id: Optional[EntityId] = None name: str entity_type: str description: Optional[str] = None @@ -47,7 +52,7 @@ class Entity(BaseModel): class AddObservationsRequest(BaseModel): """Schema for adding observations to an entity.""" - entity_id: str + entity_id: EntityId context: Optional[str] = None observations: List[Observation] @@ -67,7 +72,7 @@ class SearchNodesRequest(BaseModel): class OpenNodesRequest(BaseModel): """Request schema for open_nodes tool.""" - names: Annotated[List[str], Len(min_length=1)] + names: Annotated[List[EntityId], Len(min_length=1)] class CreateRelationsRequest(BaseModel): @@ -82,7 +87,7 @@ class CreateRelationsRequest(BaseModel): class DeleteEntityRequest(BaseModel): """Request schema for delete_entities tool.""" - entity_ids: List[str] + entity_ids: List[EntityId] class DeleteRelationsRequest(BaseModel): @@ -94,7 +99,7 @@ class DeleteRelationsRequest(BaseModel): class DeleteObservationsRequest(BaseModel): """Request schema for delete_observations tool.""" - entity_id: str + entity_id: EntityId deletions: List[Observation] @@ -118,7 +123,7 @@ class ObservationResponse(SQLAlchemyModel): class ObservationsResponse(SQLAlchemyModel): """Schema for bulk observation operation results.""" - entity_id: str + entity_id: EntityId observations: List[ObservationResponse] @@ -129,7 +134,7 @@ class RelationResponse(Relation, SQLAlchemyModel): class EntityResponse(SQLAlchemyModel): """Schema for entity data returned from the service.""" - id: str + id: EntityId name: str entity_type: str description: Optional[str] = None @@ -159,7 +164,7 @@ class OpenNodesResponse(SQLAlchemyModel): class AddObservationsResponse(SQLAlchemyModel): """Response for add_observations tool.""" - entity_id: str + entity_id: EntityId observations: List[ObservationResponse] diff --git a/src/basic_memory/utils.py b/src/basic_memory/utils.py new file mode 100644 index 00000000..8d6f3b32 --- /dev/null +++ b/src/basic_memory/utils.py @@ -0,0 +1,13 @@ +"""Utility functions for basic-memory.""" + +def normalize_entity_id(entity_id: str) -> str: + """ + Normalize an entity ID by converting to lowercase and replacing spaces with underscores. + + Args: + entity_id: Raw entity ID to normalize + + Returns: + Normalized entity ID suitable for filesystem and database use + """ + return entity_id.lower().replace(" ", "_") \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index a3840842..001d7450 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,5 @@ """Common test fixtures.""" + import tempfile from pathlib import Path from typing import AsyncGenerator @@ -7,16 +8,18 @@ import pytest_asyncio from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession, AsyncEngine from basic_memory import db +from basic_memory.config import ProjectConfig from basic_memory.db import DatabaseType -from basic_memory.repository.entity_repository import EntityRepository from basic_memory.deps import ( get_entity_service, get_observation_service, get_relation_service, - get_relation_repo, get_observation_repo, get_entity_repo + get_relation_repo, + get_observation_repo, + get_entity_repo, ) +from basic_memory.repository.entity_repository import EntityRepository from basic_memory.schemas import Entity -from basic_memory.config import ProjectConfig from basic_memory.services import MemoryService @@ -24,20 +27,25 @@ from basic_memory.services import MemoryService def anyio_backend(): return "asyncio" + @pytest_asyncio.fixture def test_config(tmp_path): """Test configuration using in-memory DB.""" config = ProjectConfig( name="test", ) - config.path=tmp_path + config.path = tmp_path return config @pytest_asyncio.fixture(scope="function") -async def engine_session_factory(test_config)-> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: +async def engine_session_factory( + test_config, +) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: """Create an async engine using in-memory SQLite database""" - async with db.engine_session_factory(project_path=test_config.path, db_type=DatabaseType.MEMORY) as (engine, session_factory): + async with db.engine_session_factory( + project_path=test_config.path, db_type=DatabaseType.MEMORY + ) as (engine, session_factory): yield engine, session_factory @@ -45,9 +53,10 @@ async def engine_session_factory(test_config)-> AsyncGenerator[tuple[AsyncEngine async def session(engine_session_factory): """Create an async session factory and yield a session""" engine, session_factory = engine_session_factory - async with session_factory() as session: + async with db.session(session_factory) as session: yield session + @pytest_asyncio.fixture async def test_project_path(): """Create a temporary project directory.""" @@ -57,62 +66,61 @@ async def test_project_path(): entities_path.mkdir(parents=True) yield project_path + @pytest_asyncio.fixture(scope="function") async def entity_repository(session: AsyncSession): """Create an EntityRepository instance""" return await get_entity_repo(session) + @pytest_asyncio.fixture(scope="function") async def observation_repository(session: AsyncSession): """Create an ObservationRepository instance""" return await get_observation_repo(session) + @pytest_asyncio.fixture(scope="function") async def relation_repository(session: AsyncSession): """Create a RelationRepository instance""" return await get_relation_repo(session) + @pytest_asyncio.fixture async def entity_service(test_project_path, entity_repository): """Fixture providing initialized EntityService.""" return await get_entity_service(test_project_path, entity_repository) + @pytest_asyncio.fixture async def relation_service(test_project_path, relation_repository): """Fixture providing initialized RelationService.""" return await get_relation_service(test_project_path, relation_repository) + @pytest_asyncio.fixture async def observation_service(test_project_path, observation_repository): """Fixture providing initialized RelationService.""" return await get_observation_service(test_project_path, observation_repository) + @pytest_asyncio.fixture -async def memory_service( - test_project_path, - entity_service, - relation_service, - observation_service -): +async def memory_service(test_project_path, entity_service, relation_service, observation_service): """Fixture providing initialized MemoryService.""" - return MemoryService( - test_project_path, - entity_service, - relation_service, - observation_service - ) + return MemoryService(test_project_path, entity_service, relation_service, observation_service) + @pytest_asyncio.fixture(scope="function") async def sample_entity(entity_repository: EntityRepository): """Create a sample entity for testing""" entity_data = { - 'id': '20240102-test-entity', - 'name': 'Test Entity', - 'entity_type': 'test', - 'description': 'A test entity', + "id": "20240102-test-entity", + "name": "Test Entity", + "entity_type": "test", + "description": "A test entity", } return await entity_repository.create(entity_data) + @pytest_asyncio.fixture async def test_entity(entity_service): """Create a test entity for reuse in tests.""" @@ -121,4 +129,3 @@ async def test_entity(entity_service): entity_type="test", # pyright: ignore [reportCallIssue] ) return await entity_service.create_entity(entity_data) - diff --git a/tests/mcp/test_mcp_server.py b/tests/mcp/test_mcp_server.py index ed9a4582..e29f897d 100644 --- a/tests/mcp/test_mcp_server.py +++ b/tests/mcp/test_mcp_server.py @@ -11,6 +11,7 @@ from basic_memory.api.app import app as fastapi_app from basic_memory.deps import get_project_config, get_engine_factory from basic_memory.mcp.server import MemoryServer, MIME_TYPE, BASIC_MEMORY_URI from basic_memory.schemas import CreateEntityResponse, SearchNodesResponse, AddObservationsResponse +from basic_memory.utils import normalize_entity_id @pytest_asyncio.fixture @@ -70,6 +71,106 @@ def test_directory_entity_data(): } +@pytest.mark.asyncio +async def test_create_single_entity(server): + """Test creating a single entity.""" + entity_data = { + "entities": [ + {"name": "SingleTest", "entity_type": "test", "observations": ["Test observation"]} + ] + } + + result = await server.handle_call_tool("create_entities", entity_data) + + # Verify response format + assert len(result) == 1 + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert result[0].resource.mimeType == MIME_TYPE + + # Verify entity creation + response = CreateEntityResponse.model_validate_json(result[0].resource.text) + assert len(response.entities) == 1 + entity = response.entities[0] + assert entity.name == "SingleTest" + assert entity.entity_type == "test" + assert len(entity.observations) == 1 + assert entity.observations[0].content == "Test observation" + assert entity.id == "test/singletest" + + # Verify entity can be found via search + search_result = await server.handle_call_tool("search_nodes", {"query": "SingleTest"}) + search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + assert len(search_response.matches) == 1 + assert search_response.matches[0].name == "SingleTest" + + +@pytest.mark.asyncio +async def test_create_multiple_entities(server): + """Test creating multiple entities in one call.""" + entity_data = { + "entities": [ + {"name": "BulkTest1", "entity_type": "test", "observations": ["First bulk test"]}, + {"name": "BulkTest2", "entity_type": "test", "observations": ["Second bulk test"]}, + {"name": "BulkTest3", "entity_type": "demo", "observations": ["Third bulk test"]}, + ] + } + + result = await server.handle_call_tool("create_entities", entity_data) + + # Verify response + assert len(result) == 1 + response = CreateEntityResponse.model_validate_json(result[0].resource.text) + + # Verify all entities were created + assert len(response.entities) == 3 + + # Check specific entities + entities = {e.name: e for e in response.entities} + assert "BulkTest1" in entities + assert "BulkTest2" in entities + assert "BulkTest3" in entities + + # Verify IDs were generated correctly + assert entities["BulkTest1"].id == "test/bulktest1" + assert entities["BulkTest2"].id == "test/bulktest2" + assert entities["BulkTest3"].id == "demo/bulktest3" + + # Verify observations were saved + assert len(entities["BulkTest1"].observations) == 1 + assert entities["BulkTest1"].observations[0].content == "First bulk test" + + # Verify entities can be found via search + search_result = await server.handle_call_tool("search_nodes", {"query": "BulkTest"}) + search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + assert len(search_response.matches) == 3 + + +@pytest.mark.asyncio +async def test_create_entity_with_all_fields(server): + """Test creating entity with all possible fields populated.""" + entity_data = { + "entities": [ + { + "name": "FullEntity", + "entity_type": "test", + "description": "A complete test entity", + "observations": ["First observation", "Second observation"], + } + ] + } + + result = await server.handle_call_tool("create_entities", entity_data) + response = CreateEntityResponse.model_validate_json(result[0].resource.text) + + entity = response.entities[0] + assert entity.name == "FullEntity" + assert entity.description == "A complete test entity" + assert len(entity.observations) == 2 + assert entity.observations[0].content == "First observation" + assert entity.observations[1].content == "Second observation" + + @pytest.mark.asyncio async def test_list_tools(server): """Test that server exposes expected tools.""" @@ -116,7 +217,7 @@ async def test_search_nodes(test_entity_data, client, server): assert result[0].resource.mimeType == MIME_TYPE # Verify search results - response = SearchNodesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] + response = SearchNodesResponse.model_validate_json(result[0].resource.text) assert len(response.matches) == 1 assert response.matches[0].name == "Test Entity" assert response.query == "Test Entity" @@ -135,7 +236,7 @@ async def test_add_observations(test_entity_data, client, server): # First create an entity create_result = await server.handle_call_tool("create_entities", test_entity_data) - create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] + create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) entity_id = create_response.entities[0].id # Add new observation @@ -162,6 +263,43 @@ async def test_add_observations(test_entity_data, client, server): assert "A new observation" in [o["content"] for o in entity["observations"]] +@pytest.mark.asyncio +async def test_create_relations(test_entity_data, client, server): + """Test creating relations between entities.""" + # Create two test entities + entity_data = { + "entities": [ + {"name": "TestEntityA", "entity_type": "test", "observations": ["Entity A"]}, + {"name": "TestEntityB", "entity_type": "test", "observations": ["Entity B"]}, + ] + } + + await server.handle_call_tool("create_entities", entity_data) + + # Create relation between them + relation_data = { + "relations": [ + { + "from_id": "test/TestEntityA", + "to_id": "test/TestEntityB", + "relation_type": "relates_to", + } + ] + } + + result = await server.handle_call_tool("create_relations", relation_data) + + # Verify through search + search_result = await server.handle_call_tool("search_nodes", {"query": "TestEntityA"}) + response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + + assert len(response.matches) == 1 + entity = response.matches[0] + assert len(entity.relations) == 1 + assert entity.relations[0].to_id == normalize_entity_id("test/TestEntityB") + assert entity.relations[0].relation_type == "relates_to" + + @pytest.mark.asyncio async def test_invalid_tool_name(server): """Test calling a non-existent tool.""" diff --git a/tests/repository/test_observation_repository.py b/tests/repository/test_observation_repository.py index 77bfb16a..ed2e010b 100644 --- a/tests/repository/test_observation_repository.py +++ b/tests/repository/test_observation_repository.py @@ -1,6 +1,8 @@ """Tests for the ObservationRepository.""" + import pytest import pytest_asyncio +import sqlalchemy from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.models import Entity, Observation @@ -12,38 +14,54 @@ async def repo(observation_repository): """Create an ObservationRepository instance""" return observation_repository + @pytest_asyncio.fixture(scope="function") async def sample_observation(repo, sample_entity: Entity): """Create a sample observation for testing""" observation_data = { - 'entity_id': sample_entity.id, - 'content': 'Test observation', - 'context': 'test-context' + "entity_id": sample_entity.id, + "content": "Test observation", + "context": "test-context", } return await repo.create(observation_data) + @pytest.mark.asyncio async def test_create_observation( - observation_repository: ObservationRepository, - sample_entity: Entity + observation_repository: ObservationRepository, sample_entity: Entity ): """Test creating a new observation""" observation_data = { - 'entity_id': sample_entity.id, - 'content': 'Test content', - 'context': 'test-context' + "entity_id": sample_entity.id, + "content": "Test content", + "context": "test-context", } observation = await observation_repository.create(observation_data) assert observation.entity_id == sample_entity.id - assert observation.content == 'Test content' + assert observation.content == "Test content" assert observation.id is not None # Should be auto-generated + +@pytest.mark.asyncio +async def test_create_observation_entity_does_not_exist( + observation_repository: ObservationRepository, sample_entity: Entity +): + """Test creating a new observation""" + observation_data = { + "entity_id": "does-not-exist", + "content": "Test content", + "context": "test-context", + } + with pytest.raises(sqlalchemy.exc.IntegrityError): + await observation_repository.create(observation_data) + + @pytest.mark.asyncio async def test_find_by_entity( - observation_repository: ObservationRepository, - sample_observation: Observation, - sample_entity: Entity + observation_repository: ObservationRepository, + sample_observation: Observation, + sample_entity: Entity, ): """Test finding observations by entity""" observations = await observation_repository.find_by_entity(sample_entity.id) @@ -51,13 +69,13 @@ async def test_find_by_entity( assert observations[0].id == sample_observation.id assert observations[0].content == sample_observation.content + @pytest.mark.asyncio async def test_find_by_context( - observation_repository: ObservationRepository, - sample_observation: Observation + observation_repository: ObservationRepository, sample_observation: Observation ): """Test finding observations by context""" - observations = await observation_repository.find_by_context('test-context') + observations = await observation_repository.find_by_context("test-context") assert len(observations) == 1 assert observations[0].id == sample_observation.id assert observations[0].content == sample_observation.content @@ -68,10 +86,7 @@ async def test_delete_observations(session: AsyncSession, repo): """Test deleting observations by entity_id.""" # Create test entity entity = Entity( - id="test/test_entity", - name="test_entity", - entity_type="test", - description="Test entity" + id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" ) session.add(entity) await session.flush() @@ -90,15 +105,13 @@ async def test_delete_observations(session: AsyncSession, repo): remaining = await repo.find_by_entity(entity.id) assert len(remaining) == 0 + @pytest.mark.asyncio async def test_delete_observation_by_id(session: AsyncSession, repo): """Test deleting a single observation by its ID.""" # Create test entity entity = Entity( - id="test/test_entity", - name="test_entity", - entity_type="test", - description="Test entity" + id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" ) session.add(entity) await session.flush() @@ -116,15 +129,13 @@ async def test_delete_observation_by_id(session: AsyncSession, repo): remaining = await repo.find_by_id(obs.id) assert remaining is None + @pytest.mark.asyncio async def test_delete_observation_by_content(session: AsyncSession, repo): """Test deleting observations by content.""" # Create test entity entity = Entity( - id="test/test_entity", - name="test_entity", - entity_type="test", - description="Test entity" + id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" ) session.add(entity) await session.flush() @@ -135,7 +146,6 @@ async def test_delete_observation_by_content(session: AsyncSession, repo): session.add_all([obs1, obs2]) await session.flush() - # Test deletion by content deleted = await repo.delete_by_fields(content="Delete this observation") assert deleted is True diff --git a/tests/repository/test_relation_repository.py b/tests/repository/test_relation_repository.py index efb07463..5532624f 100644 --- a/tests/repository/test_relation_repository.py +++ b/tests/repository/test_relation_repository.py @@ -1,11 +1,13 @@ """Tests for the RelationRepository.""" + import pytest import pytest_asyncio +import sqlalchemy + from basic_memory.models import Entity, Relation from basic_memory.repository.relation_repository import RelationRepository - @pytest_asyncio.fixture async def relation_repo(session): """Create a RelationRepository with test DB session.""" @@ -19,7 +21,7 @@ async def source_entity(session): id="source/test_entity", name="test_source", entity_type="source", - description="Source entity" + description="Source entity", ) session.add(entity) await session.flush() @@ -33,7 +35,7 @@ async def target_entity(session): id="target/test_entity", name="test_target", entity_type="target", - description="Target entity" + description="Target entity", ) session.add(entity) await session.flush() @@ -44,181 +46,181 @@ async def target_entity(session): async def test_relations(session, source_entity, target_entity): """Create test relations.""" relations = [ - Relation( - from_id=source_entity.id, - to_id=target_entity.id, - relation_type="connects_to" - ), - Relation( - from_id=source_entity.id, - to_id=target_entity.id, - relation_type="depends_on" - ) + Relation(from_id=source_entity.id, to_id=target_entity.id, relation_type="connects_to"), + Relation(from_id=source_entity.id, to_id=target_entity.id, relation_type="depends_on"), ] session.add_all(relations) await session.flush() return relations + @pytest_asyncio.fixture(scope="function") 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', - 'references': '' + "id": "20240102-related", + "name": "Related Entity", + "entity_type": "test", + "description": "A related test entity", + "references": "", } return await entity_repository.create(entity_data) + @pytest_asyncio.fixture(scope="function") async def sample_relation( - relation_repository: RelationRepository, - sample_entity: Entity, - related_entity: Entity + relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity ): """Create a sample relation for testing""" relation_data = { - 'from_id': sample_entity.id, - 'to_id': related_entity.id, - 'relation_type': 'test_relation', - 'context': 'test-context' + "from_id": sample_entity.id, + "to_id": related_entity.id, + "relation_type": "test_relation", + "context": "test-context", } return await relation_repository.create(relation_data) + @pytest_asyncio.fixture(scope="function") async def multiple_relations( - relation_repository: RelationRepository, - sample_entity: Entity, - related_entity: Entity + relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity ): """Create multiple relations for testing""" relations_data = [ { - 'from_id': sample_entity.id, - 'to_id': related_entity.id, - 'relation_type': 'relation_one', - 'context': 'context_one' + "from_id": sample_entity.id, + "to_id": related_entity.id, + "relation_type": "relation_one", + "context": "context_one", }, { - 'from_id': sample_entity.id, - 'to_id': related_entity.id, - 'relation_type': 'relation_two', - 'context': 'context_two' + "from_id": sample_entity.id, + "to_id": related_entity.id, + "relation_type": "relation_two", + "context": "context_two", }, { - 'from_id': related_entity.id, - 'to_id': sample_entity.id, - 'relation_type': 'relation_one', - 'context': 'context_three' - } + "from_id": related_entity.id, + "to_id": sample_entity.id, + "relation_type": "relation_one", + "context": "context_three", + }, ] return [await relation_repository.create(data) for data in relations_data] + @pytest.mark.asyncio async def test_create_relation( - relation_repository: RelationRepository, - sample_entity: Entity, - related_entity: Entity + relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity ): """Test creating a new relation""" relation_data = { - 'from_id': sample_entity.id, - 'to_id': related_entity.id, - 'relation_type': 'test_relation', - 'context': 'test-context' + "from_id": sample_entity.id, + "to_id": related_entity.id, + "relation_type": "test_relation", + "context": "test-context", } relation = await relation_repository.create(relation_data) assert relation.from_id == sample_entity.id assert relation.to_id == related_entity.id - assert relation.relation_type == 'test_relation' + assert relation.relation_type == "test_relation" assert relation.id is not None # Should be auto-generated + +@pytest.mark.asyncio +async def test_create_relation_entity_does_not_exist( + relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity +): + """Test creating a new relation""" + relation_data = { + "from_id": "not_exist", + "to_id": related_entity.id, + "relation_type": "test_relation", + "context": "test-context", + } + with pytest.raises(sqlalchemy.exc.IntegrityError): + await relation_repository.create(relation_data) + + @pytest.mark.asyncio async def test_find_by_entities( - relation_repository: RelationRepository, - sample_relation: Relation, - sample_entity: Entity, - related_entity: Entity + relation_repository: RelationRepository, + sample_relation: Relation, + sample_entity: Entity, + related_entity: Entity, ): """Test finding relations between specific entities""" - relations = await relation_repository.find_by_entities( - sample_entity.id, - related_entity.id - ) + relations = await relation_repository.find_by_entities(sample_entity.id, related_entity.id) assert len(relations) == 1 assert relations[0].id == sample_relation.id assert relations[0].relation_type == sample_relation.relation_type + @pytest.mark.asyncio -async def test_find_by_type( - relation_repository: RelationRepository, - sample_relation: Relation -): +async def test_find_by_type(relation_repository: RelationRepository, sample_relation: Relation): """Test finding relations by type""" - relations = await relation_repository.find_by_type('test_relation') + relations = await relation_repository.find_by_type("test_relation") assert len(relations) == 1 assert relations[0].id == sample_relation.id + @pytest.mark.asyncio async def test_delete_by_fields_single_field( - relation_repository: RelationRepository, - multiple_relations: list[Relation] + relation_repository: RelationRepository, multiple_relations: list[Relation] ): """Test deleting relations by a single field.""" # Delete all relations of type 'relation_one' - result = await relation_repository.delete_by_fields(relation_type='relation_one') # pyright: ignore [reportArgumentType] + result = await relation_repository.delete_by_fields(relation_type="relation_one") # pyright: ignore [reportArgumentType] assert result is True # Verify deletion - remaining = await relation_repository.find_by_type('relation_one') + remaining = await relation_repository.find_by_type("relation_one") assert len(remaining) == 0 # Other relations should still exist - others = await relation_repository.find_by_type('relation_two') + others = await relation_repository.find_by_type("relation_two") assert len(others) == 1 + @pytest.mark.asyncio async def test_delete_by_fields_multiple_fields( - relation_repository: RelationRepository, - multiple_relations: list[Relation], - sample_entity: Entity, - related_entity: Entity + relation_repository: RelationRepository, + multiple_relations: list[Relation], + sample_entity: Entity, + related_entity: Entity, ): """Test deleting relations by multiple fields.""" # Delete specific relation matching both from_id and relation_type result = await relation_repository.delete_by_fields( - from_id=sample_entity.id, # pyright: ignore [reportArgumentType] - relation_type='relation_one' # pyright: ignore [reportArgumentType] + from_id=sample_entity.id, # pyright: ignore [reportArgumentType] + relation_type="relation_one", # pyright: ignore [reportArgumentType] ) assert result is True # Verify correct relation was deleted - remaining = await relation_repository.find_by_entities( - sample_entity.id, - related_entity.id - ) + remaining = await relation_repository.find_by_entities(sample_entity.id, related_entity.id) assert len(remaining) == 1 # Only relation_two should remain - assert remaining[0].relation_type == 'relation_two' + assert remaining[0].relation_type == "relation_two" + @pytest.mark.asyncio async def test_delete_by_fields_no_match( - relation_repository: RelationRepository, - multiple_relations: list[Relation] + relation_repository: RelationRepository, multiple_relations: list[Relation] ): """Test delete_by_fields when no relations match.""" result = await relation_repository.delete_by_fields( - relation_type='nonexistent_type' # pyright: ignore [reportArgumentType] + relation_type="nonexistent_type" # pyright: ignore [reportArgumentType] ) assert result is False + @pytest.mark.asyncio async def test_delete_by_fields_all_fields( - relation_repository: RelationRepository, - multiple_relations: list[Relation], - sample_entity: Entity, - related_entity: Entity + relation_repository: RelationRepository, + multiple_relations: list[Relation], + sample_entity: Entity, + related_entity: Entity, ): """Test deleting relation by matching all fields.""" # Get first relation's data @@ -226,10 +228,10 @@ async def test_delete_by_fields_all_fields( # Delete using all fields result = await relation_repository.delete_by_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] + 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 @@ -239,15 +241,14 @@ async def test_delete_by_fields_all_fields( assert remaining[0].context != relation.context - @pytest.mark.asyncio async def test_delete_relation_by_id(relation_repo, test_relations): """Test deleting a relation by ID.""" relation = test_relations[0] - + result = await relation_repo.delete(relation.id) assert result is True - + # Verify deletion remaining = await relation_repo.find_one( relation_repo.select(Relation).filter(Relation.id == relation.id) @@ -260,25 +261,24 @@ async def test_delete_relations_by_type(relation_repo, test_relations): """Test deleting relations by type.""" result = await relation_repo.delete_by_fields(relation_type="connects_to") assert result is True - + # Verify specific type was deleted remaining = await relation_repo.find_by_type("connects_to") assert len(remaining) == 0 - + # Verify other type still exists others = await relation_repo.find_by_type("depends_on") assert len(others) == 1 @pytest.mark.asyncio -async def test_delete_relations_by_entities(relation_repo, test_relations, source_entity, target_entity): +async def test_delete_relations_by_entities( + relation_repo, test_relations, source_entity, target_entity +): """Test deleting relations between specific entities.""" - result = await relation_repo.delete_by_fields( - from_id=source_entity.id, - to_id=target_entity.id - ) + result = await relation_repo.delete_by_fields(from_id=source_entity.id, to_id=target_entity.id) assert result is True - + # Verify all relations between entities were deleted remaining = await relation_repo.find_by_entities(source_entity.id, target_entity.id) assert len(remaining) == 0 @@ -288,4 +288,4 @@ async def test_delete_relations_by_entities(relation_repo, test_relations, sourc async def test_delete_nonexistent_relation(relation_repo): """Test deleting a relation that doesn't exist.""" result = await relation_repo.delete_by_fields(relation_type="nonexistent") - assert result is False \ No newline at end of file + assert result is False