learningcircuit--local-deep-research/tests/database/test_encrypted_database_orm.py

#!/usr/bin/env python3
"""Test ORM operations with encrypted user databases."""

import sys
import tempfile
import uuid
from datetime import datetime, timezone
from pathlib import Path

import pytest

sys.path.insert(
    0,
    str(Path(__file__).parent.parent.parent.resolve()),
)

from local_deep_research.database.encrypted_db import DatabaseManager
from local_deep_research.database.models import (
    APIKey,
    Report,
    ResearchHistory,
    ResearchLog,
    ResearchResource,
    ResearchTask,
    SearchQuery,
    SearchResult,
    UserSettings,
)


class TestEncryptedDatabaseORM:
    """Test ORM operations in encrypted user databases."""

    @pytest.fixture
    def temp_data_dir(self):
        """Create a temporary directory for test databases."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield Path(tmpdir)

    @pytest.fixture
    def db_manager(self, temp_data_dir, monkeypatch):
        """Create a database manager with temporary directory."""
        monkeypatch.setattr(
            "local_deep_research.database.encrypted_db.get_data_directory",
            lambda: temp_data_dir,
        )
        manager = DatabaseManager()
        yield manager
        # Cleanup
        for username in list(manager.connections.keys()):
            manager.close_user_database(username)

    @pytest.fixture
    def test_user_session(self, db_manager):
        """Create a test user with encrypted database and return session."""
        username = "testuser"
        password = "TestPassword123!"

        # Create user database
        db_manager.create_user_database(username, password)

        # Get session
        session = db_manager.get_session(username)
        yield session

        # Cleanup
        session.close()
        db_manager.close_user_database(username)

    def test_user_settings_crud(self, test_user_session):
        """Test CRUD operations on UserSettings."""
        session = test_user_session

        # Create settings
        setting1 = UserSettings(
            key="llm.provider",
            value={"provider": "openai", "model": "gpt-4"},
            category="llm",
        )
        setting2 = UserSettings(
            key="search.engine",
            value={"engine": "duckduckgo", "safe_search": True},
            category="search",
        )

        session.add_all([setting1, setting2])
        session.commit()

        # Read settings
        llm_setting = (
            session.query(UserSettings).filter_by(key="llm.provider").first()
        )
        assert llm_setting is not None
        assert llm_setting.value["provider"] == "openai"
        assert llm_setting.category == "llm"

        # Update setting
        llm_setting.value = {"provider": "anthropic", "model": "claude-3"}
        session.commit()

        # Verify update
        updated = (
            session.query(UserSettings).filter_by(key="llm.provider").first()
        )
        assert updated.value["provider"] == "anthropic"

        # Delete setting
        session.delete(setting2)
        session.commit()

        # Verify deletion
        remaining = session.query(UserSettings).count()
        assert remaining == 1

    def test_api_keys_encryption(self, test_user_session):
        """Test API key storage and retrieval."""
        session = test_user_session

        # Store API keys
        openai_key = APIKey(
            provider="openai",
            key="sk-test123456789",  # This will be encrypted in the database
            is_active=True,
        )
        anthropic_key = APIKey(
            provider="anthropic", key="sk-ant-test987654321", is_active=False
        )

        session.add_all([openai_key, anthropic_key])
        session.commit()

        # Retrieve and verify
        stored_key = session.query(APIKey).filter_by(provider="openai").first()
        assert stored_key is not None
        assert stored_key.key == "sk-test123456789"  # Should decrypt properly
        assert stored_key.is_active is True

        # Query by active status
        active_keys = session.query(APIKey).filter_by(is_active=True).all()
        assert len(active_keys) == 1
        assert active_keys[0].provider == "openai"

    def test_research_history_with_resources(self, test_user_session):
        """Test ResearchHistory with related ResearchResource objects."""
        session = test_user_session

        # Create research history
        research_id = str(uuid.uuid4())
        research = ResearchHistory(
            id=research_id,
            query="quantum computing applications",
            mode="detailed",
            status="completed",
            created_at=datetime.now(timezone.utc).isoformat(),
            completed_at=datetime.now(timezone.utc).isoformat(),
            progress=100,
            research_meta={
                "model": "gpt-4",
                "temperature": 0.7,
                "iterations": 3,
            },
        )
        session.add(research)
        session.commit()

        # Add resources
        resources = [
            ResearchResource(
                research_id=research_id,
                title="Quantum Computing in Drug Discovery",
                url="https://example.com/quantum-drug",
                content_preview="Recent advances in quantum computing...",
                source_type="article",
                resource_metadata={"author": "Dr. Smith", "year": 2024},
                created_at=datetime.now(timezone.utc).isoformat(),
            ),
            ResearchResource(
                research_id=research_id,
                title="IBM Quantum Network",
                url="https://example.com/ibm-quantum",
                content_preview="IBM's quantum computing initiative...",
                source_type="web",
                resource_metadata={"company": "IBM", "relevance": 0.95},
                created_at=datetime.now(timezone.utc).isoformat(),
            ),
        ]

        session.add_all(resources)
        session.commit()

        # Query with relationship
        stored_research = (
            session.query(ResearchHistory).filter_by(id=research_id).first()
        )
        assert stored_research is not None
        assert len(stored_research.resources) == 2

        # Query resources directly
        quantum_resources = (
            session.query(ResearchResource)
            .filter(ResearchResource.title.contains("Quantum"))
            .all()
        )
        assert len(quantum_resources) == 2

    def test_research_task_workflow(self, test_user_session):
        """Test complete research task workflow with queries and results."""
        session = test_user_session

        # Create research task
        task = ResearchTask(
            title="AI Safety Research",
            description="Comprehensive research on AI alignment and safety",
            status="in_progress",
        )
        session.add(task)
        session.commit()

        # Add search queries
        query1 = SearchQuery(
            research_task_id=task.id,
            query="AI alignment problem solutions",
            search_engine="google",
            search_type="academic",
            status="completed",
        )
        query2 = SearchQuery(
            research_task_id=task.id,
            query="AI safety research organizations",
            search_engine="duckduckgo",
            search_type="web",
            status="completed",
        )

        session.add_all([query1, query2])
        session.commit()

        # Add search results
        results = [
            SearchResult(
                research_task_id=task.id,
                search_query_id=query1.id,
                title="Concrete Problems in AI Safety",
                url="https://example.com/ai-safety-paper",
                snippet="A comprehensive survey of AI safety challenges...",
                relevance_score=0.98,
                position=1,
            ),
            SearchResult(
                research_task_id=task.id,
                search_query_id=query2.id,
                title="Center for AI Safety",
                url="https://example.com/cais",
                snippet="Leading research organization focused on AI safety...",
                relevance_score=0.95,
                position=1,
            ),
        ]

        session.add_all(results)
        session.commit()

        # Add report
        report = Report(
            research_task_id=task.id,
            title="AI Safety Research Report",
            content="# AI Safety Research\n\n## Executive Summary\n...",
            format="markdown",
            is_draft=False,
            word_count=1500,
            section_count=5,
        )

        session.add(report)
        session.commit()

        # Update task status
        task.status = "completed"
        task.completed_at = datetime.now(timezone.utc)
        session.commit()

        # Verify relationships
        completed_task = (
            session.query(ResearchTask).filter_by(id=task.id).first()
        )
        assert completed_task.status == "completed"
        assert len(completed_task.searches) == 2
        assert len(completed_task.results) == 2
        assert len(completed_task.reports) == 1

        # Query high relevance results
        high_relevance = (
            session.query(SearchResult)
            .filter(SearchResult.relevance_score > 0.9)
            .all()
        )
        assert len(high_relevance) == 2

    def test_research_logs(self, test_user_session):
        """Test research logging functionality."""
        session = test_user_session

        # ResearchLog.research_id is a String(36) FK to
        # research_history.id (UUID), not research.id (Integer).
        # Production writes UUIDs via log_utils, so the test must too.
        research = ResearchHistory(
            id=str(uuid.uuid4()),
            query="Test research for logging",
            mode="detailed",
            status="in_progress",
            created_at=datetime.now(timezone.utc).isoformat(),
        )
        session.add(research)
        session.commit()

        # Add various log entries
        logs = [
            ResearchLog(
                research_id=research.id,
                timestamp=datetime.now(timezone.utc),
                level="INFO",
                message="Starting research process",
                module="research_service",
                function="start_research",
                line_no=100,
            ),
            ResearchLog(
                research_id=research.id,
                timestamp=datetime.now(timezone.utc),
                level="DEBUG",
                message="Executing search query",
                module="search_engine",
                function="search",
                line_no=250,
            ),
            ResearchLog(
                research_id=research.id,
                timestamp=datetime.now(timezone.utc),
                level="ERROR",
                message="API rate limit exceeded",
                module="api_client",
                function="make_request",
                line_no=75,
            ),
        ]

        session.add_all(logs)
        session.commit()

        # Query logs by level
        error_logs = (
            session.query(ResearchLog)
            .filter_by(research_id=research.id, level="ERROR")
            .all()
        )
        assert len(error_logs) == 1
        assert "rate limit" in error_logs[0].message

        # Query all logs for research
        all_logs = (
            session.query(ResearchLog)
            .filter_by(research_id=research.id)
            .order_by(ResearchLog.timestamp)
            .all()
        )
        assert len(all_logs) == 3

    def test_database_integrity(self, db_manager):
        """Test database integrity checks."""
        username = "integrity_test_user"
        password = "IntegrityTest123!"

        # Create and open database
        db_manager.create_user_database(username, password)

        # Check integrity
        integrity_ok = db_manager.check_database_integrity(username)
        assert integrity_ok is True

        # Cleanup
        db_manager.close_user_database(username)

    def test_concurrent_users(self, db_manager):
        """Test multiple users with separate encrypted databases."""
        users = [
            ("alice", "AlicePass123!"),
            ("bob", "BobPass456!"),
            ("charlie", "CharliePass789!"),
        ]

        # Create databases for each user
        for username, password in users:
            engine = db_manager.create_user_database(username, password)
            assert engine is not None

            # Add a setting specific to each user
            session = db_manager.get_session(username)
            setting = UserSettings(
                key="user.theme",
                value={"theme": f"{username}_theme"},
                category="ui",
            )
            session.add(setting)
            session.commit()
            session.close()

        # Verify each user has their own data
        for username, password in users:
            session = db_manager.get_session(username)
            setting = (
                session.query(UserSettings).filter_by(key="user.theme").first()
            )
            assert setting is not None
            assert setting.value["theme"] == f"{username}_theme"
            session.close()

        # Cleanup
        for username, _ in users:
            db_manager.close_user_database(username)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])