#!/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"])