129 lines
4.3 KiB
Python
129 lines
4.3 KiB
Python
"""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
|
|
|
|
|
|
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 == 42
|
|
extraction_repo.session.commit.assert_called_once()
|
|
extraction_repo.session.refresh.assert_called_once_with(extraction)
|