sdb2-backend/tests/repositories/test_extraction.py

"""Tests for extraction repository."""

from unittest.mock import AsyncMock, Mock

import pytest
from sqlmodel.ext.asyncio.session import AsyncSession

from app.models.extraction import Extraction
from app.repositories.extraction import ExtractionRepository

# Constants
TEST_SOUND_ID = 42


class TestExtractionRepository:
    """Test extraction repository."""

    @pytest.fixture
    def mock_session(self):
        """Create a mock session."""
        return Mock(spec=AsyncSession)

    @pytest.fixture
    def extraction_repo(self, mock_session):
        """Create an extraction repository with mock session."""
        return ExtractionRepository(mock_session)

    def test_init(self, extraction_repo):
        """Test repository initialization."""
        assert extraction_repo.session is not None

    @pytest.mark.asyncio
    async def test_create_extraction(self, extraction_repo):
        """Test creating an extraction."""
        extraction_data = {
            "url": "https://www.youtube.com/watch?v=test",
            "user_id": 1,
            "service": "youtube",
            "service_id": "test123",
            "title": "Test Video",
            "status": "pending",
        }

        # Mock the session operations
        extraction_repo.session.add = Mock()
        extraction_repo.session.commit = AsyncMock()
        extraction_repo.session.refresh = AsyncMock()

        result = await extraction_repo.create(extraction_data)

        # Verify the result has the expected attributes
        assert result.url == extraction_data["url"]
        assert result.user_id == extraction_data["user_id"]
        assert result.service == extraction_data["service"]
        assert result.service_id == extraction_data["service_id"]
        assert result.title == extraction_data["title"]
        assert result.status == extraction_data["status"]

        # Verify session methods were called
        extraction_repo.session.add.assert_called_once()
        extraction_repo.session.commit.assert_called_once()
        extraction_repo.session.refresh.assert_called_once()

    @pytest.mark.asyncio
    async def test_get_by_service_and_id(self, extraction_repo):
        """Test getting extraction by service and service_id."""
        mock_result = Mock()
        mock_result.first.return_value = Extraction(
            id=1,
            service="youtube",
            service_id="test123",
            url="https://www.youtube.com/watch?v=test",
            user_id=1,
            status="pending",
        )

        extraction_repo.session.exec = AsyncMock(return_value=mock_result)

        result = await extraction_repo.get_by_service_and_id("youtube", "test123")

        assert result is not None
        assert result.service == "youtube"
        assert result.service_id == "test123"
        extraction_repo.session.exec.assert_called_once()

    @pytest.mark.asyncio
    async def test_get_pending_extractions(self, extraction_repo):
        """Test getting pending extractions."""
        mock_extraction = Extraction(
            id=1,
            service="youtube",
            service_id="test123",
            url="https://www.youtube.com/watch?v=test",
            user_id=1,
            status="pending",
        )

        mock_result = Mock()
        mock_result.all.return_value = [mock_extraction]

        extraction_repo.session.exec = AsyncMock(return_value=mock_result)

        result = await extraction_repo.get_pending_extractions()

        assert len(result) == 1
        assert result[0].status == "pending"
        extraction_repo.session.exec.assert_called_once()

    @pytest.mark.asyncio
    async def test_update_extraction(self, extraction_repo):
        """Test updating an extraction."""
        extraction = Extraction(
            id=1,
            service="youtube",
            service_id="test123",
            url="https://www.youtube.com/watch?v=test",
            user_id=1,
            status="pending",
        )

        update_data = {"status": "completed", "sound_id": 42}

        extraction_repo.session.commit = AsyncMock()
        extraction_repo.session.refresh = AsyncMock()

        result = await extraction_repo.update(extraction, update_data)

        assert result.status == "completed"
        assert result.sound_id == TEST_SOUND_ID
        extraction_repo.session.commit.assert_called_once()
        extraction_repo.session.refresh.assert_called_once_with(extraction)

    @pytest.mark.asyncio
    async def test_get_by_status(self, extraction_repo):
        """Test getting extractions by status."""
        mock_extractions = [
            Extraction(
                id=1,
                service="youtube",
                service_id="test123",
                url="https://www.youtube.com/watch?v=test1",
                user_id=1,
                status="processing",
            ),
            Extraction(
                id=2,
                service="youtube",
                service_id="test456",
                url="https://www.youtube.com/watch?v=test2",
                user_id=1,
                status="processing",
            ),
        ]

        mock_result = Mock()
        mock_result.all.return_value = mock_extractions

        extraction_repo.session.exec = AsyncMock(return_value=mock_result)

        result = await extraction_repo.get_by_status("processing")

        assert len(result) == 2
        assert all(extraction.status == "processing" for extraction in result)
        extraction_repo.session.exec.assert_called_once()