pig.py/tests/unit/test_database.py

"""
Tests for database module.

Tests SQLAlchemy models, database operations, and queries.
"""

import pytest
from unittest.mock import Mock, AsyncMock, patch, MagicMock
import tempfile
import shutil
import os
import json
from datetime import datetime, timezone

from database import (
    Database,
    CrawlJob,
    Media,
    MediaSource,
    Page,
    BackfillJob,
    SCORE_SCREENSHOT,
    SCORE_OG_IMAGE,
    SCORE_THUMBNAIL,
    SCORE_FULL_RES,
    OVER_9000,
)


class TestScoreConstants:
    """Test scoring constants."""

    def test_score_hierarchy(self):
        """Test that scores follow expected hierarchy."""
        assert SCORE_SCREENSHOT < SCORE_OG_IMAGE
        assert SCORE_OG_IMAGE < SCORE_THUMBNAIL
        assert SCORE_THUMBNAIL < SCORE_FULL_RES

    def test_score_values(self):
        """Test specific score values."""
        assert SCORE_SCREENSHOT == 1
        assert SCORE_OG_IMAGE == 3
        assert SCORE_THUMBNAIL == 5
        assert SCORE_FULL_RES == 10

    def test_over_9000(self):
        """Test the limit constant."""
        assert OVER_9000 == 9000


class TestDatabaseInitialization:
    """Test Database class initialization."""

    def setup_method(self):
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, "test.db")

    def teardown_method(self):
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    @pytest.mark.asyncio
    async def test_database_init(self):
        """Test database initialization creates tables."""
        db = Database(self.db_path)
        await db.init()

        assert db._initialized is True
        assert db._engine is not None
        assert db._session_factory is not None
        assert os.path.exists(self.db_path)

        await db.close()

    @pytest.mark.asyncio
    async def test_database_double_init(self):
        """Test that double initialization is safe."""
        db = Database(self.db_path)
        await db.init()
        await db.init()  # Should not raise

        assert db._initialized is True
        await db.close()

    @pytest.mark.asyncio
    async def test_database_session(self):
        """Test getting a session."""
        db = Database(self.db_path)
        await db.init()

        session = db.session()
        assert session is not None

        await db.close()


class TestCrawlJobOperations:
    """Test crawl job database operations."""

    def setup_method(self):
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, "test.db")

    def teardown_method(self):
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    @pytest.mark.asyncio
    async def test_create_crawl_job(self):
        """Test creating a crawl job."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(
            target_uri="https://example.com",
            keywords=["test", "demo"],
            mode="images",
            depth=10,
            max_pages=100,
            fast=False,
            screenshots=True,
        )

        assert job_id > 0

        await db.close()

    @pytest.mark.asyncio
    async def test_create_crawl_job_defaults(self):
        """Test creating crawl job with defaults."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        assert job_id > 0

        await db.close()

    @pytest.mark.asyncio
    async def test_update_crawl_job_stats(self):
        """Test updating crawl job stats."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")
        stats = {"pages_crawled": 10, "media_found": 50}

        await db.update_crawl_job_stats(job_id, stats)

        # Verify stats were saved
        job = await db.get_crawl_job(job_id)
        assert job is not None
        saved_stats = json.loads(job['stats'])
        assert saved_stats['pages_crawled'] == 10

        await db.close()

    @pytest.mark.asyncio
    async def test_complete_crawl_job(self):
        """Test completing a crawl job."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")
        stats = {"pages_crawled": 100, "media_found": 500}

        await db.complete_crawl_job(job_id, stats)

        job = await db.get_crawl_job(job_id)
        assert job['status'] == 'completed'
        assert job['completed_at'] is not None

        await db.close()

    @pytest.mark.asyncio
    async def test_pause_crawl_job(self):
        """Test pausing a crawl job."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")
        stats = {"pages_crawled": 50}

        await db.pause_crawl_job(job_id, stats)

        job = await db.get_crawl_job(job_id)
        assert job['status'] == 'paused'

        await db.close()

    @pytest.mark.asyncio
    async def test_fail_crawl_job(self):
        """Test failing a crawl job."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        await db.fail_crawl_job(job_id, "Connection timeout")

        job = await db.get_crawl_job(job_id)
        assert job['status'] == 'failed'
        assert job['error'] == 'Connection timeout'

        await db.close()

    @pytest.mark.asyncio
    async def test_delete_crawl_job(self):
        """Test deleting a crawl job."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        result = await db.delete_crawl_job(job_id)

        assert result['deleted'] is True
        assert result['target_uri'] == "https://example.com"

        await db.close()

    @pytest.mark.asyncio
    async def test_delete_nonexistent_job(self):
        """Test deleting a job that doesn't exist."""
        db = Database(self.db_path)
        await db.init()

        result = await db.delete_crawl_job(99999)

        assert result['deleted'] is False

        await db.close()


class TestMediaOperations:
    """Test media database operations."""

    def setup_method(self):
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, "test.db")

    def teardown_method(self):
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    @pytest.mark.asyncio
    async def test_create_media_record(self):
        """Test creating a media record."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        await db.create_media_record(
            md5_hash="d41d8cd98f00b204e9800998ecf8427e",
            media_uri="https://example.com/image.jpg",
            page_uri="https://example.com/",
            crawl_job_id=job_id,
            media_type="image",
            mime_type="image/jpeg",
            file_size=12345,
            page_title="Example Page",
            alt_text="A sample image",
            score=SCORE_FULL_RES,
        )

        # Verify media was created
        media = await db.get_media_by_hash("d41d8cd98f00b204e9800998ecf8427e")
        assert media is not None
        assert media['media_type'] == 'image'

        await db.close()

    @pytest.mark.asyncio
    async def test_create_duplicate_media(self):
        """Test that duplicate media is deduplicated."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        # Create first media
        await db.create_media_record(
            md5_hash="d41d8cd98f00b204e9800998ecf8427e",
            media_uri="https://example.com/image.jpg",
            page_uri="https://example.com/page1",
            crawl_job_id=job_id,
            media_type="image",
        )

        # Create same media from different page
        await db.create_media_record(
            md5_hash="d41d8cd98f00b204e9800998ecf8427e",
            media_uri="https://example.com/image.jpg",
            page_uri="https://example.com/page2",
            crawl_job_id=job_id,
            media_type="image",
        )

        # Should still have only one media record
        media = await db.get_media_by_hash("d41d8cd98f00b204e9800998ecf8427e")
        assert media is not None

        # But should have two sources
        sources = await db.get_media_sources("d41d8cd98f00b204e9800998ecf8427e")
        # Note: depends on unique constraint behavior
        assert len(sources) >= 1

        await db.close()


class TestPageOperations:
    """Test page database operations."""

    def setup_method(self):
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, "test.db")

    def teardown_method(self):
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    @pytest.mark.asyncio
    async def test_create_page(self):
        """Test creating a page record."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        await db.store_page(
            uri="https://example.com/about",
            title="About Page",
            content="This is the about page content.",
            crawl_job_id=job_id,
        )

        # Verify page was created
        page = await db.get_page_by_uri("https://example.com/about")
        assert page is not None
        assert page['title'] == 'About Page'

        await db.close()

    @pytest.mark.asyncio
    async def test_backfill_page_hashes(self):
        """Test backfilling page URI hashes."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        await db.store_page(
            uri="https://example.com/test",
            title="Test",
            content="Test content",
            crawl_job_id=job_id,
        )

        count = await db.backfill_page_hashes()
        # Should have backfilled at least one
        assert count >= 0

        await db.close()


class TestSearchOperations:
    """Test search functionality."""

    def setup_method(self):
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, "test.db")

    def teardown_method(self):
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    @pytest.mark.asyncio
    async def test_search_media(self):
        """Test searching for media."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        # Create media with searchable text
        await db.create_media_record(
            md5_hash="abc123def456",
            media_uri="https://example.com/sunset.jpg",
            page_uri="https://example.com/",
            crawl_job_id=job_id,
            media_type="image",
            alt_text="Beautiful sunset over the ocean",
            searchable_text="sunset ocean beautiful landscape",
        )

        results = await db.search_media_advanced(q="sunset", limit=10)
        # Result depends on search implementation
        assert isinstance(results, list)

        await db.close()

    @pytest.mark.asyncio
    async def test_get_stats(self):
        """Test getting database stats."""
        db = Database(self.db_path)
        await db.init()

        stats = await db.get_stats()

        assert 'media_count' in stats or 'images' in stats or isinstance(stats, dict)

        await db.close()

    @pytest.mark.asyncio
    async def test_get_recent_media(self):
        """Test getting recent media."""
        db = Database(self.db_path)
        await db.init()

        job_id = await db.create_crawl_job(target_uri="https://example.com")

        # Create some media
        for i in range(5):
            await db.create_media_record(
                md5_hash=f"hash{i}{'0' * 27}",
                media_uri=f"https://example.com/image{i}.jpg",
                page_uri="https://example.com/",
                crawl_job_id=job_id,
                media_type="image",
            )

        recent = await db.get_recent_media(limit=3)
        assert len(recent) <= 3

        await db.close()


class TestCrawlJobModel:
    """Test CrawlJob model."""

    def test_crawl_job_table_name(self):
        """Test CrawlJob table name."""
        assert CrawlJob.__tablename__ == 'crawl_jobs'

    def test_crawl_job_columns(self):
        """Test CrawlJob has required columns."""
        columns = [c.name for c in CrawlJob.__table__.columns]
        assert 'id' in columns
        assert 'target_uri' in columns
        assert 'keywords' in columns
        assert 'mode' in columns
        assert 'status' in columns
        assert 'started_at' in columns
        assert 'completed_at' in columns
        assert 'stats' in columns


class TestMediaModel:
    """Test Media model."""

    def test_media_table_name(self):
        """Test Media table name."""
        assert Media.__tablename__ == 'media'

    def test_media_primary_key(self):
        """Test Media primary key is md5_hash."""
        pk_columns = [c.name for c in Media.__table__.primary_key.columns]
        assert 'md5_hash' in pk_columns

    def test_media_columns(self):
        """Test Media has required columns."""
        columns = [c.name for c in Media.__table__.columns]
        assert 'md5_hash' in columns
        assert 'media_type' in columns
        assert 'mime_type' in columns
        assert 'file_size' in columns
        assert 'score' in columns
        assert 'first_seen_at' in columns


class TestMediaSourceModel:
    """Test MediaSource model."""

    def test_media_source_table_name(self):
        """Test MediaSource table name."""
        assert MediaSource.__tablename__ == 'media_sources'

    def test_media_source_columns(self):
        """Test MediaSource has required columns."""
        columns = [c.name for c in MediaSource.__table__.columns]
        assert 'md5_hash' in columns
        assert 'media_uri' in columns
        assert 'page_uri' in columns
        assert 'page_title' in columns
        assert 'alt_text' in columns
        assert 'discovered_at' in columns


class TestPageModel:
    """Test Page model."""

    def test_page_table_name(self):
        """Test Page table name."""
        assert Page.__tablename__ == 'pages'

    def test_page_columns(self):
        """Test Page has required columns."""
        columns = [c.name for c in Page.__table__.columns]
        assert 'uri' in columns
        assert 'uri_hash' in columns
        assert 'title' in columns
        assert 'content' in columns


class TestBackfillJobModel:
    """Test BackfillJob model."""

    def test_backfill_job_table_name(self):
        """Test BackfillJob table name."""
        assert BackfillJob.__tablename__ == 'backfill_jobs'

    def test_backfill_job_columns(self):
        """Test BackfillJob has required columns."""
        columns = [c.name for c in BackfillJob.__table__.columns]
        assert 'job_type' in columns
        assert 'status' in columns
        assert 'total_records' in columns
        assert 'processed_records' in columns


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