#!/usr/bin/env python3 """Test that ORM conversions work correctly (no more raw SQL).""" import sys from pathlib import Path sys.path.insert( 0, str(Path(__file__).parent.parent.parent.resolve()), ) import pytest from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from local_deep_research.database.models import ( Base, ResearchHistory, ResearchLog, ResearchResource, ) @pytest.fixture def test_db(): """Create a test database in memory.""" engine = create_engine("sqlite:///:memory:") Base.metadata.create_all(engine) Session = sessionmaker(bind=engine) session = Session() yield session session.close() engine.dispose() def test_research_history_orm_queries(test_db): """Test ResearchHistory ORM queries work correctly.""" import uuid # Create test data with UUID research = ResearchHistory( id=str(uuid.uuid4()), query="Test quantum computing", mode="deep", status="completed", created_at="2024-01-01T00:00:00", progress=100, research_meta={"model": "gpt-4", "iterations": 5}, ) test_db.add(research) test_db.commit() # Test querying by ID found = test_db.query(ResearchHistory).filter_by(id=research.id).first() assert found is not None assert found.query == "Test quantum computing" # Test querying by status completed = ( test_db.query(ResearchHistory).filter_by(status="completed").all() ) assert len(completed) == 1 assert completed[0].id == research.id # Test ordering ordered = ( test_db.query(ResearchHistory) .order_by(ResearchHistory.created_at.desc()) .all() ) assert len(ordered) == 1 print("āœ“ ResearchHistory ORM queries work correctly") def test_research_resource_orm_operations(test_db): """Test ResearchResource ORM operations.""" import uuid # Create a research entry first with UUID research = ResearchHistory( id=str(uuid.uuid4()), query="Test", mode="quick", status="completed", created_at="2024-01-01T00:00:00", ) test_db.add(research) test_db.commit() # Add resources resource1 = ResearchResource( research_id=research.id, title="Resource 1", url="https://example.com/1", content_preview="Preview 1", source_type="web", resource_metadata={"relevance": 0.9}, created_at="2024-01-01T12:00:00", ) resource2 = ResearchResource( research_id=research.id, title="Resource 2", url="https://example.com/2", content_preview="Preview 2", source_type="pdf", created_at="2024-01-01T12:05:00", ) test_db.add_all([resource1, resource2]) test_db.commit() # Query resources for research resources = ( test_db.query(ResearchResource) .filter_by(research_id=research.id) .order_by(ResearchResource.id.asc()) .all() ) assert len(resources) == 2 assert resources[0].title == "Resource 1" assert resources[1].title == "Resource 2" assert resources[0].resource_metadata == {"relevance": 0.9} # Test deletion test_db.delete(resource1) test_db.commit() remaining = ( test_db.query(ResearchResource) .filter_by(research_id=research.id) .count() ) assert remaining == 1 print("āœ“ ResearchResource ORM operations work correctly") def test_research_log_orm_queries(test_db): """Test ResearchLog ORM queries.""" from datetime import datetime, timezone # First create a Research entry (not ResearchHistory) from local_deep_research.database.models import ( Research, ResearchMode, ResearchStatus, ) research = Research( query="Test research", status=ResearchStatus.IN_PROGRESS, mode=ResearchMode.QUICK, ) test_db.add(research) test_db.commit() # Add logs with all required fields log1 = ResearchLog( research_id=research.id, timestamp=datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc), message="Starting research", module="research_service", function="start_research", line_no=100, level="INFO", ) log2 = ResearchLog( research_id=research.id, timestamp=datetime(2024, 1, 1, 0, 1, 0, tzinfo=timezone.utc), message="Search completed", module="search_engine", function="search", line_no=250, level="INFO", ) test_db.add_all([log1, log2]) test_db.commit() # Query logs logs = ( test_db.query(ResearchLog) .filter_by(research_id=research.id) .order_by(ResearchLog.timestamp.asc()) .all() ) assert len(logs) == 2 assert logs[0].message == "Starting research" assert logs[1].level == "INFO" # Query by module search_logs = ( test_db.query(ResearchLog) .filter_by(research_id=research.id, module="search_engine") .all() ) assert len(search_logs) == 1 assert search_logs[0].message == "Search completed" print("āœ“ ResearchLog ORM queries work correctly") if __name__ == "__main__": # Create in-memory database for testing engine = create_engine("sqlite:///:memory:") Base.metadata.create_all(engine) Session = sessionmaker(bind=engine) session = Session() try: test_research_history_orm_queries(session) test_research_resource_orm_operations(session) test_research_log_orm_queries(session) print("\nāœ… All ORM conversion tests passed!") finally: session.close() engine.dispose()