local-deep-research/tests/test_database_initialization.py

"""Test database initialization module"""

import os
import tempfile
from pathlib import Path
import pytest
from sqlalchemy import create_engine, inspect
from sqlalchemy.orm import sessionmaker

from local_deep_research.database.initialize import (
    initialize_database,
    check_database_schema,
)
from local_deep_research.database.models import (
    Setting,
)


class TestDatabaseInitialization:
    """Test the centralized database initialization"""

    @pytest.fixture
    def temp_db(self):
        """Create a temporary database for testing"""
        with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
            db_path = tmp.name

        engine = create_engine(f"sqlite:///{db_path}")
        yield engine, db_path

        # Cleanup
        engine.dispose()
        if Path(db_path).exists():
            os.unlink(db_path)

    def test_initialize_core_tables(self, temp_db):
        """Test that core tables are created correctly"""
        engine, db_path = temp_db

        # Initialize database
        Session = sessionmaker(bind=engine)
        with Session() as session:
            initialize_database(engine, session)

        # Check that core tables exist
        inspector = inspect(engine)
        tables = inspector.get_table_names()

        # Verify essential tables (users is auth-only, not in per-user DBs)
        expected_core_tables = [
            "settings",
            "research",
            "research_history",
            "journals",
            "app_logs",  # Correct table name
            "queued_researches",  # Correct table name
            "search_cache",
            "token_usage",
            "research_ratings",
        ]

        for table in expected_core_tables:
            assert table in tables, f"Table '{table}' should exist"

    def test_initialize_with_news_tables(self, temp_db):
        """Test that news tables are created when requested"""
        engine, db_path = temp_db

        # Initialize database with news tables
        Session = sessionmaker(bind=engine)
        with Session() as session:
            initialize_database(engine, session)

        # Check schema
        schema_info = check_database_schema(engine)

        # News tables might not be available if the news module isn't imported
        # Just check that initialization doesn't fail
        assert len(schema_info["tables"]) > 0
        assert len(schema_info["missing_tables"]) == 0

    def test_idempotent_initialization(self, temp_db):
        """Test that initialization is idempotent (can be run multiple times)"""
        engine, db_path = temp_db
        Session = sessionmaker(bind=engine)

        # Initialize twice
        with Session() as session:
            initialize_database(engine, session)

        # Get initial table count
        inspector = inspect(engine)
        initial_tables = set(inspector.get_table_names())

        # Initialize again
        with Session() as session:
            initialize_database(engine, session)

        # Check that no duplicate tables were created
        final_tables = set(inspector.get_table_names())
        assert initial_tables == final_tables

    def test_check_database_schema(self, temp_db):
        """Test the schema checking function"""
        engine, db_path = temp_db

        # Check schema before initialization
        schema_info = check_database_schema(engine)
        assert len(schema_info["tables"]) == 0
        assert len(schema_info["missing_tables"]) > 0

        # Initialize database
        Session = sessionmaker(bind=engine)
        with Session() as session:
            initialize_database(engine, session)

        # Check schema after initialization
        schema_info = check_database_schema(engine)
        assert len(schema_info["tables"]) > 0
        assert len(schema_info["missing_tables"]) == 0

    def test_settings_initialization(self, temp_db):
        """Test that settings can be initialized"""
        engine, db_path = temp_db

        # Initialize database with settings
        Session = sessionmaker(bind=engine)
        with Session() as session:
            initialize_database(engine, session)

            # Check if settings table is queryable
            count = session.query(Setting).count()
            assert count >= 0  # Should not raise an error

    def test_partial_table_creation(self, temp_db):
        """Test that initialization completes even with existing tables"""
        engine, db_path = temp_db

        # Create only one table manually (Setting has no foreign keys)
        from local_deep_research.database.models import Setting

        Setting.__table__.create(engine)

        # Verify only that table exists
        inspector = inspect(engine)
        initial_tables = inspector.get_table_names()
        assert "settings" in initial_tables
        assert len(initial_tables) == 1  # Only settings table

        # Initialize database (should create all missing tables)
        Session = sessionmaker(bind=engine)
        with Session() as session:
            initialize_database(engine, session)

        # Verify many more tables now exist (need fresh inspector)
        final_inspector = inspect(engine)
        final_tables = final_inspector.get_table_names()
        assert "settings" in final_tables  # Original table still there
        assert len(final_tables) > 20  # Many more tables created

        # Verify some key tables
        assert "research" in final_tables
        assert "journals" in final_tables
        assert "app_logs" in final_tables
        assert "token_usage" in final_tables