basicmachines-co-basic-memory/tests/conftest.py

"""Common test fixtures."""

import os
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from textwrap import dedent
from typing import AsyncGenerator

import pytest
import pytest_asyncio
from alembic import command
from alembic.config import Config
from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.pool import NullPool
from testcontainers.postgres import PostgresContainer

from basic_memory import db
from basic_memory.config import (
    ProjectConfig,
    ProjectEntry,
    BasicMemoryConfig,
    ConfigManager,
    DatabaseBackend,
)
from basic_memory.db import DatabaseType
from basic_memory.markdown import EntityParser
from basic_memory.markdown.markdown_processor import MarkdownProcessor
from basic_memory.models import Base
from basic_memory.models.knowledge import Entity
from basic_memory.models.project import Project
from basic_memory.repository.entity_repository import EntityRepository
from basic_memory.repository.observation_repository import ObservationRepository
from basic_memory.repository.project_repository import ProjectRepository
from basic_memory.repository.relation_repository import RelationRepository
from basic_memory.schemas.base import Entity as EntitySchema
from basic_memory.services import (
    EntityService,
    ProjectService,
)
from basic_memory.services.directory_service import DirectoryService
from basic_memory.services.file_service import FileService
from basic_memory.services.link_resolver import LinkResolver
from basic_memory.services.search_service import SearchService
from basic_memory.sync.sync_service import SyncService
from basic_memory.sync.watch_service import WatchService


# =============================================================================
# Database Backend Selection (env var approach)
# =============================================================================
# By default, tests run against SQLite.
# Set BASIC_MEMORY_TEST_POSTGRES=1 to run against Postgres (uses testcontainers).
# This allows running sqlite/postgres tests in parallel in CI.


@pytest.fixture(scope="session")
def db_backend():
    """Determine database backend from environment variable.

    Default: sqlite
    Set BASIC_MEMORY_TEST_POSTGRES=1 to use postgres
    """
    if os.environ.get("BASIC_MEMORY_TEST_POSTGRES", "").lower() in ("1", "true", "yes"):
        return "postgres"
    return "sqlite"


@pytest.fixture(scope="session")
def postgres_container(db_backend):
    """Session-scoped Postgres container for tests.

    Uses testcontainers to spin up a real Postgres instance in Docker.
    The container is started once per test session and shared across all tests.
    Only starts if db_backend is "postgres".
    """
    if db_backend != "postgres" or _configured_postgres_sync_url():
        yield None
        return

    # Use pgvector image so CREATE EXTENSION vector succeeds in search repository
    with PostgresContainer("pgvector/pgvector:pg16") as postgres:
        yield postgres


POSTGRES_EPHEMERAL_TABLES = [
    "search_vector_embeddings",
    "search_vector_index",
]


def _configured_postgres_sync_url() -> str | None:
    """Prefer an externally managed Postgres server when CI provides one."""
    configured_url = os.environ.get("BASIC_MEMORY_TEST_POSTGRES_URL") or os.environ.get(
        "POSTGRES_TEST_URL"
    )
    if not configured_url:
        return None

    return (
        configured_url.replace("postgresql+asyncpg://", "postgresql+psycopg2://", 1)
        .replace("postgresql://", "postgresql+psycopg2://", 1)
        .replace("postgres://", "postgresql+psycopg2://", 1)
    )


def _postgres_alembic_config(async_url: str) -> Config:
    """Build Alembic config for stamping the shared Postgres test schema."""
    alembic_dir = Path(db.__file__).parent / "alembic"
    cfg = Config()
    cfg.set_main_option("script_location", str(alembic_dir))
    cfg.set_main_option(
        "file_template",
        "%%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s",
    )
    cfg.set_main_option("timezone", "UTC")
    cfg.set_main_option("revision_environment", "false")
    cfg.set_main_option("sqlalchemy.url", async_url)
    return cfg


def _postgres_reset_tables() -> list[str]:
    """Resolve the current ORM table set at reset time.

    Some tests declare models after conftest import, so the list must stay dynamic.
    """
    return [table.name for table in Base.metadata.sorted_tables] + [
        "search_index",
        "search_vector_chunks",
    ]


def _resolve_postgres_sync_url(postgres_container) -> str:
    """Use CI's shared service when configured, otherwise fall back to testcontainers."""
    configured_url = _configured_postgres_sync_url()
    if configured_url:
        return configured_url
    assert postgres_container is not None
    return postgres_container.get_connection_url()


async def _reset_postgres_test_schema(engine: AsyncEngine, async_url: str) -> None:
    """Restore the shared Postgres schema to a clean baseline."""
    from basic_memory.models.search import (
        CREATE_POSTGRES_SEARCH_INDEX_FTS,
        CREATE_POSTGRES_SEARCH_INDEX_METADATA,
        CREATE_POSTGRES_SEARCH_INDEX_PERMALINK,
        CREATE_POSTGRES_SEARCH_INDEX_TABLE,
        CREATE_POSTGRES_SEARCH_VECTOR_CHUNKS_INDEX,
        CREATE_POSTGRES_SEARCH_VECTOR_CHUNKS_TABLE,
    )

    async with engine.begin() as conn:
        # Trigger: several tests intentionally drop or stub search tables to exercise recovery code.
        # Why: TRUNCATE is much cheaper than drop_all/create_all, but it only works when the schema exists.
        # Outcome: we recreate any missing core tables once, then clear rows for deterministic test setup.
        await conn.run_sync(Base.metadata.create_all)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_TABLE)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_FTS)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_METADATA)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_PERMALINK)
        await conn.execute(CREATE_POSTGRES_SEARCH_VECTOR_CHUNKS_TABLE)
        await conn.execute(CREATE_POSTGRES_SEARCH_VECTOR_CHUNKS_INDEX)

        for table_name in POSTGRES_EPHEMERAL_TABLES:
            await conn.execute(text(f"DROP TABLE IF EXISTS {table_name} CASCADE"))

        await conn.execute(
            text(f"TRUNCATE TABLE {', '.join(_postgres_reset_tables())} RESTART IDENTITY CASCADE")
        )

        alembic_version_exists = (
            await conn.execute(text("SELECT to_regclass('public.alembic_version')"))
        ).scalar() is not None

    if not alembic_version_exists:
        command.stamp(_postgres_alembic_config(async_url), "head")


@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def postgres_engine(
    db_backend, postgres_container
) -> AsyncGenerator[AsyncEngine | None, None]:
    """Create the shared Postgres engine once per test session."""
    if db_backend != "postgres":
        yield None
        return

    sync_url = _resolve_postgres_sync_url(postgres_container)
    async_url = sync_url.replace("postgresql+psycopg2", "postgresql+asyncpg")
    engine = create_async_engine(
        async_url,
        echo=False,
        poolclass=NullPool,
    )

    try:
        yield engine
    finally:
        await engine.dispose()


@pytest.fixture
def anyio_backend():
    return "asyncio"


@pytest.fixture(autouse=True)
def suppress_logfire_no_config_warning(monkeypatch) -> None:
    """Keep tests focused on behavior instead of Logfire bootstrap warnings."""
    monkeypatch.setenv("LOGFIRE_IGNORE_NO_CONFIG", "1")


@pytest.fixture(autouse=True)
def isolate_routing_env(monkeypatch) -> None:
    """Prevent command-routing env flags from leaking across tests."""
    monkeypatch.delenv("BASIC_MEMORY_FORCE_LOCAL", raising=False)
    monkeypatch.delenv("BASIC_MEMORY_FORCE_CLOUD", raising=False)
    monkeypatch.delenv("BASIC_MEMORY_EXPLICIT_ROUTING", raising=False)


@pytest.fixture(autouse=True)
def isolate_data_dir_env(monkeypatch) -> None:
    """Keep host data-dir env vars from leaking into tests.

    Why: GitHub Actions Ubuntu runners set ``XDG_CONFIG_HOME=/home/runner/.config``,
    and ``resolve_data_dir()`` honors it ahead of ``Path.home() / ".basic-memory"``.
    Without clearing it, tests that monkeypatch HOME still see the host XDG path
    and assertions against the tmp home directory fail.
    """
    monkeypatch.delenv("BASIC_MEMORY_CONFIG_DIR", raising=False)
    monkeypatch.delenv("XDG_CONFIG_HOME", raising=False)


@pytest_asyncio.fixture(autouse=True)
async def cleanup_global_db_after_test() -> AsyncGenerator[None, None]:
    """Close any module-level DB engine created outside fixture ownership."""
    yield

    # Trigger: a test exercises production fallback routing instead of the
    # per-test engine fixture.
    # Why: that path stores an engine in basic_memory.db module state, and
    # a later fixture can overwrite the reference before it is disposed.
    # Outcome: close straggler aiosqlite worker threads before the loop closes.
    await db.shutdown_db()


@pytest.fixture
def project_root() -> Path:
    return Path(__file__).parent.parent


@pytest.fixture
def config_home(tmp_path, monkeypatch) -> Path:
    # Patch HOME environment variable for the duration of the test
    monkeypatch.setenv("HOME", str(tmp_path))
    # On Windows, also set USERPROFILE
    if os.name == "nt":
        monkeypatch.setenv("USERPROFILE", str(tmp_path))
    # Set BASIC_MEMORY_HOME to the test directory
    monkeypatch.setenv("BASIC_MEMORY_HOME", str(tmp_path / "basic-memory"))
    return tmp_path


@pytest.fixture(scope="function")
def app_config(config_home, db_backend, postgres_container, monkeypatch) -> BasicMemoryConfig:
    """Create test app configuration for the appropriate backend."""
    projects = {"test-project": ProjectEntry(path=str(config_home))}

    # Set backend based on parameterized db_backend fixture
    if db_backend == "postgres":
        backend = DatabaseBackend.POSTGRES
        # Trigger: CI jobs can provide a shared Postgres service instead of per-session containers.
        # Why: reusing one pgvector-enabled server avoids Docker startup churn on every job.
        # Outcome: local runs keep using testcontainers, while CI injects a stable service URL.
        sync_url = _resolve_postgres_sync_url(postgres_container)
        database_url = sync_url.replace("postgresql+psycopg2", "postgresql+asyncpg")
    else:
        backend = DatabaseBackend.SQLITE
        database_url = None

    app_config = BasicMemoryConfig(
        env="test",
        projects=projects,
        default_project="test-project",
        update_permalinks_on_move=True,
        database_backend=backend,
        database_url=database_url,
    )

    return app_config


@pytest.fixture
def config_manager(app_config: BasicMemoryConfig, config_home: Path, monkeypatch) -> ConfigManager:
    # Invalidate config cache to ensure clean state for each test
    from basic_memory import config as config_module

    config_module._CONFIG_CACHE = None
    config_module._CONFIG_MTIME = None
    config_module._CONFIG_SIZE = None

    # Create a new ConfigManager that uses the test home directory
    config_manager = ConfigManager()
    # Update its paths to use the test directory
    config_manager.config_dir = config_home / ".basic-memory"
    config_manager.config_file = config_manager.config_dir / "config.json"
    config_manager.config_dir.mkdir(parents=True, exist_ok=True)

    # Ensure the config file is written to disk
    config_manager.save_config(app_config)
    return config_manager


@pytest.fixture(scope="function")
def project_config(test_project):
    """Create test project configuration."""

    project_config = ProjectConfig(
        name=test_project.name,
        home=Path(test_project.path),
    )

    return project_config


@dataclass
class TestConfig:
    config_home: Path
    project_config: ProjectConfig
    app_config: BasicMemoryConfig
    config_manager: ConfigManager


@pytest.fixture
def test_config(config_home, project_config, app_config, config_manager) -> TestConfig:
    """All test configuration fixtures"""
    return TestConfig(config_home, project_config, app_config, config_manager)


@pytest_asyncio.fixture(scope="function")
async def engine_factory(
    app_config,
    config_manager,
    db_backend,
    postgres_container,
    postgres_engine,
) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
    """Engine factory for SQLite or Postgres tests.

    Uses parameterized db_backend fixture to run tests against both backends.
    """
    from basic_memory.models.search import (
        CREATE_SEARCH_INDEX,
        CREATE_SQLITE_SEARCH_VECTOR_CHUNKS,
        CREATE_SQLITE_SEARCH_VECTOR_CHUNKS_PROJECT_ENTITY,
        CREATE_SQLITE_SEARCH_VECTOR_CHUNKS_UNIQUE,
    )

    if db_backend == "postgres":
        assert postgres_engine is not None
        sync_url = _resolve_postgres_sync_url(postgres_container)
        async_url = sync_url.replace("postgresql+psycopg2", "postgresql+asyncpg")

        # Trigger: Basic Memory sync/indexing paths can catch database errors,
        # roll back, and continue through new scoped sessions.
        # Why: binding every scoped session to one savepoint-backed connection can
        # leave that connection in PendingRollbackError state for later sessions.
        # Outcome: keep the session-scoped engine, but use normal Postgres
        # transaction semantics and reset rows/schema between tests.
        await _reset_postgres_test_schema(postgres_engine, async_url)

        session_maker = async_sessionmaker(
            bind=postgres_engine,
            class_=AsyncSession,
            expire_on_commit=False,
            autoflush=False,
        )

        # Important: wire the engine/session into the global db module state.
        # Some codepaths (e.g. app initialization / MCP lifespan) call db.get_or_create_db(),
        # which would otherwise create a separate engine and run migrations, conflicting with
        # our test-created schema (and causing DuplicateTableError).
        db._engine = postgres_engine
        db._session_maker = session_maker

        try:
            yield postgres_engine, session_maker
        finally:
            if db._engine is postgres_engine:
                db._engine = None
            if db._session_maker is session_maker:
                db._session_maker = None
    else:
        # SQLite mode
        db_type = DatabaseType.MEMORY
        async with db.engine_session_factory(db_path=app_config.database_path, db_type=db_type) as (
            engine,
            session_maker,
        ):
            # Create all tables via ORM, then add search_index via FTS5 DDL
            async with engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
                await conn.execute(CREATE_SEARCH_INDEX)
                await conn.execute(CREATE_SQLITE_SEARCH_VECTOR_CHUNKS)
                await conn.execute(CREATE_SQLITE_SEARCH_VECTOR_CHUNKS_PROJECT_ENTITY)
                await conn.execute(CREATE_SQLITE_SEARCH_VECTOR_CHUNKS_UNIQUE)

            # Yield after setup is complete
            yield engine, session_maker


@pytest_asyncio.fixture
async def session_maker(engine_factory) -> async_sessionmaker[AsyncSession]:
    """Get session maker for tests."""
    _, session_maker = engine_factory
    return session_maker


## Repositories


@pytest_asyncio.fixture(scope="function")
async def entity_repository(
    session_maker: async_sessionmaker[AsyncSession], test_project: Project
) -> EntityRepository:
    """Create an EntityRepository instance with project context."""
    return EntityRepository(session_maker, project_id=test_project.id)


@pytest_asyncio.fixture(scope="function")
async def observation_repository(
    session_maker: async_sessionmaker[AsyncSession], test_project: Project
) -> ObservationRepository:
    """Create an ObservationRepository instance with project context."""
    return ObservationRepository(session_maker, project_id=test_project.id)


@pytest_asyncio.fixture(scope="function")
async def relation_repository(
    session_maker: async_sessionmaker[AsyncSession], test_project: Project
) -> RelationRepository:
    """Create a RelationRepository instance with project context."""
    return RelationRepository(session_maker, project_id=test_project.id)


@pytest_asyncio.fixture(scope="function")
async def project_repository(
    session_maker: async_sessionmaker[AsyncSession],
) -> ProjectRepository:
    """Create a ProjectRepository instance."""
    return ProjectRepository(session_maker)


@pytest_asyncio.fixture(scope="function")
async def test_project(config_home, engine_factory) -> Project:
    """Create a test project to be used as context for other repositories."""
    project_data = {
        "name": "test-project",
        "description": "Project used as context for tests",
        "path": str(config_home),
        "is_active": True,
        "is_default": True,  # Explicitly set as the default project (for cli operations)
    }
    engine, session_maker = engine_factory
    project_repository = ProjectRepository(session_maker)
    project = await project_repository.create(project_data)
    return project


## Services


@pytest_asyncio.fixture
async def entity_service(
    entity_repository: EntityRepository,
    observation_repository: ObservationRepository,
    relation_repository: RelationRepository,
    entity_parser: EntityParser,
    file_service: FileService,
    link_resolver: LinkResolver,
    app_config: BasicMemoryConfig,
) -> EntityService:
    """Create EntityService."""
    return EntityService(
        entity_parser=entity_parser,
        entity_repository=entity_repository,
        observation_repository=observation_repository,
        relation_repository=relation_repository,
        file_service=file_service,
        link_resolver=link_resolver,
        app_config=app_config,
    )


@pytest.fixture
def file_service(
    project_config: ProjectConfig, markdown_processor: MarkdownProcessor
) -> FileService:
    """Create FileService instance."""
    return FileService(project_config.home, markdown_processor)


@pytest.fixture
def markdown_processor(entity_parser: EntityParser) -> MarkdownProcessor:
    """Create writer instance."""
    return MarkdownProcessor(entity_parser)


@pytest.fixture
def link_resolver(entity_repository: EntityRepository, search_service: SearchService):
    """Create parser instance."""
    return LinkResolver(entity_repository, search_service)


@pytest.fixture
def entity_parser(project_config):
    """Create parser instance."""
    return EntityParser(project_config.home)


@pytest_asyncio.fixture
async def sync_service(
    app_config: BasicMemoryConfig,
    entity_service: EntityService,
    entity_parser: EntityParser,
    project_repository: ProjectRepository,
    entity_repository: EntityRepository,
    relation_repository: RelationRepository,
    search_service: SearchService,
    file_service: FileService,
) -> SyncService:
    """Create sync service for testing."""
    return SyncService(
        app_config=app_config,
        entity_service=entity_service,
        project_repository=project_repository,
        entity_repository=entity_repository,
        relation_repository=relation_repository,
        entity_parser=entity_parser,
        search_service=search_service,
        file_service=file_service,
    )


@pytest_asyncio.fixture
async def directory_service(entity_repository, project_config) -> DirectoryService:
    """Create directory service for testing."""
    return DirectoryService(
        entity_repository=entity_repository,
    )


@pytest_asyncio.fixture
async def search_repository(session_maker, test_project: Project, app_config: BasicMemoryConfig):
    """Create backend-appropriate SearchRepository instance with project context"""
    from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
    from basic_memory.repository.postgres_search_repository import PostgresSearchRepository

    if app_config.database_backend == DatabaseBackend.POSTGRES:
        return PostgresSearchRepository(
            session_maker,
            project_id=test_project.id,
            app_config=app_config,
        )
    else:
        return SQLiteSearchRepository(
            session_maker,
            project_id=test_project.id,
            app_config=app_config,
        )


@pytest_asyncio.fixture
async def search_service(
    search_repository,
    entity_repository: EntityRepository,
    file_service: FileService,
) -> SearchService:
    """Create and initialize search service"""
    service = SearchService(search_repository, entity_repository, file_service)
    await service.init_search_index()
    return service


@pytest_asyncio.fixture(scope="function")
async def sample_entity(entity_repository: EntityRepository) -> Entity:
    """Create a sample entity for testing."""
    entity_data = {
        "project_id": entity_repository.project_id,
        "title": "Test Entity",
        "note_type": "test",
        "permalink": "test/test-entity",
        "file_path": "test/test_entity.md",
        "content_type": "text/markdown",
        "created_at": datetime.now(timezone.utc),
        "updated_at": datetime.now(timezone.utc),
    }
    return await entity_repository.create(entity_data)


@pytest_asyncio.fixture
async def project_service(
    project_repository: ProjectRepository,
    file_service: FileService,
) -> ProjectService:
    """Create ProjectService with repository and file service for directory operations."""
    return ProjectService(repository=project_repository, file_service=file_service)


@pytest_asyncio.fixture
async def full_entity(sample_entity, entity_repository, file_service, entity_service) -> Entity:
    """Create a search test entity."""

    # Create test entity
    entity, created = await entity_service.create_or_update_entity(
        EntitySchema(
            title="Search_Entity",
            directory="test",
            note_type="test",
            content=dedent("""
                ## Observations
                - [tech] Tech note
                - [design] Design note

                ## Relations
                - out1 [[Test Entity]]
                - out2 [[Test Entity]]
                """),
        )
    )
    return entity


@pytest_asyncio.fixture
async def test_graph(
    entity_repository,
    relation_repository,
    observation_repository,
    search_service,
    file_service,
    entity_service,
):
    """Create a test knowledge graph with entities, relations and observations."""

    # Create some test entities in reverse order so they will be linked
    deeper, _ = await entity_service.create_or_update_entity(
        EntitySchema(
            title="Deeper Entity",
            note_type="deeper",
            directory="test",
            content=dedent("""
                # Deeper Entity
                """),
        )
    )

    deep, _ = await entity_service.create_or_update_entity(
        EntitySchema(
            title="Deep Entity",
            note_type="deep",
            directory="test",
            content=dedent("""
                # Deep Entity
                - deeper_connection [[Deeper Entity]]
                """),
        )
    )

    connected_2, _ = await entity_service.create_or_update_entity(
        EntitySchema(
            title="Connected Entity 2",
            note_type="test",
            directory="test",
            content=dedent("""
                # Connected Entity 2
                - deep_connection [[Deep Entity]]
                """),
        )
    )

    connected_1, _ = await entity_service.create_or_update_entity(
        EntitySchema(
            title="Connected Entity 1",
            note_type="test",
            directory="test",
            content=dedent("""
                # Connected Entity 1
                - [note] Connected 1 note
                - connected_to [[Connected Entity 2]]
                """),
        )
    )

    root, _ = await entity_service.create_or_update_entity(
        EntitySchema(
            title="Root",
            note_type="test",
            directory="test",
            content=dedent("""
                # Root Entity
                - [note] Root note 1
                - [tech] Root tech note
                - connects_to [[Connected Entity 1]]
                """),
        )
    )

    # get latest
    entities = await entity_repository.find_all()
    relations = await relation_repository.find_all()

    # Index everything for search
    for entity in entities:
        await search_service.index_entity(entity)

    return {
        "root": root,
        "connected1": connected_1,
        "connected2": connected_2,
        "deep": deep,
        "observations": [e.observations for e in entities],
        "relations": relations,
    }


@pytest.fixture
def watch_service(app_config: BasicMemoryConfig, project_repository, sync_service) -> WatchService:
    """Create WatchService with injected sync_service factory.

    The sync_service_factory allows tests to use the fixture-provided sync_service
    instead of the production get_sync_service() which creates its own db connection.
    """

    async def sync_service_factory(project):
        """Return the test fixture's sync_service regardless of project."""
        return sync_service

    return WatchService(
        app_config=app_config,
        project_repository=project_repository,
        sync_service_factory=sync_service_factory,
    )


@pytest.fixture
def test_files(project_config, project_root) -> dict[str, Path]:
    """Copy test files into the project directory.

    Returns a dict mapping file names to their paths in the project dir.
    """
    # Source files relative to tests directory
    source_files = {
        "pdf": Path(project_root / "tests/Non-MarkdownFileSupport.pdf"),
        "image": Path(project_root / "tests/Screenshot.png"),
    }

    # Create copies in temp project directory
    project_files = {}
    for name, src_path in source_files.items():
        # Read source file
        content = src_path.read_bytes()

        # Create destination path and ensure parent dirs exist
        dest_path = project_config.home / src_path.name
        dest_path.parent.mkdir(parents=True, exist_ok=True)

        # Write file
        dest_path.write_bytes(content)
        project_files[name] = dest_path

    return project_files


@pytest_asyncio.fixture
async def synced_files(sync_service, project_config, test_files):
    # Initial sync - should create forward reference
    await sync_service.sync(project_config.home)
    return test_files