feat: add ArtifactRepository with batch DB operations
Add optimized database operations for artifact storage: - Atomic upserts using ON CONFLICT for artifact creation - Batch inserts for dependencies to eliminate N+1 queries - Joined queries for cached URL lookups - All methods include comprehensive unit tests
This commit is contained in:
167
backend/tests/unit/test_db_utils.py
Normal file
167
backend/tests/unit/test_db_utils.py
Normal file
@@ -0,0 +1,167 @@
|
||||
"""Tests for database utility functions."""
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
class TestArtifactRepository:
|
||||
"""Tests for ArtifactRepository."""
|
||||
|
||||
def test_batch_dependency_values_formatting(self):
|
||||
"""batch_upsert_dependencies should format values correctly."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
deps = [
|
||||
("_pypi", "numpy", ">=1.21.0"),
|
||||
("_pypi", "requests", "*"),
|
||||
("myproject", "mylib", "==1.0.0"),
|
||||
]
|
||||
|
||||
values = ArtifactRepository._format_dependency_values("abc123", deps)
|
||||
|
||||
assert len(values) == 3
|
||||
assert values[0] == {
|
||||
"artifact_id": "abc123",
|
||||
"dependency_project": "_pypi",
|
||||
"dependency_package": "numpy",
|
||||
"version_constraint": ">=1.21.0",
|
||||
}
|
||||
assert values[2]["dependency_project"] == "myproject"
|
||||
|
||||
def test_empty_dependencies_returns_empty_list(self):
|
||||
"""Empty dependency list should return empty values."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
values = ArtifactRepository._format_dependency_values("abc123", [])
|
||||
|
||||
assert values == []
|
||||
|
||||
def test_format_dependency_values_preserves_special_characters(self):
|
||||
"""Version constraints with special characters should be preserved."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
deps = [
|
||||
("_pypi", "package-name", ">=1.0.0,<2.0.0"),
|
||||
("_pypi", "another_pkg", "~=1.4.2"),
|
||||
]
|
||||
|
||||
values = ArtifactRepository._format_dependency_values("hash123", deps)
|
||||
|
||||
assert values[0]["version_constraint"] == ">=1.0.0,<2.0.0"
|
||||
assert values[1]["version_constraint"] == "~=1.4.2"
|
||||
|
||||
def test_batch_upsert_dependencies_returns_zero_for_empty(self):
|
||||
"""batch_upsert_dependencies should return 0 for empty list without DB call."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
mock_db = MagicMock()
|
||||
repo = ArtifactRepository(mock_db)
|
||||
|
||||
result = repo.batch_upsert_dependencies("abc123", [])
|
||||
|
||||
assert result == 0
|
||||
# Verify no DB operations were performed
|
||||
mock_db.execute.assert_not_called()
|
||||
|
||||
def test_get_or_create_artifact_builds_correct_statement(self):
|
||||
"""get_or_create_artifact should use ON CONFLICT DO UPDATE."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
from backend.app.models import Artifact
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_artifact = MagicMock()
|
||||
mock_artifact.ref_count = 1
|
||||
mock_result.scalar_one.return_value = mock_artifact
|
||||
mock_db.execute.return_value = mock_result
|
||||
|
||||
repo = ArtifactRepository(mock_db)
|
||||
artifact, created = repo.get_or_create_artifact(
|
||||
sha256="abc123def456",
|
||||
size=1024,
|
||||
filename="test.whl",
|
||||
content_type="application/zip",
|
||||
)
|
||||
|
||||
assert mock_db.execute.called
|
||||
assert created is True
|
||||
assert artifact == mock_artifact
|
||||
|
||||
def test_get_or_create_artifact_existing_not_created(self):
|
||||
"""get_or_create_artifact should return created=False for existing artifact."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_artifact = MagicMock()
|
||||
mock_artifact.ref_count = 5 # Existing artifact with ref_count > 1
|
||||
mock_result.scalar_one.return_value = mock_artifact
|
||||
mock_db.execute.return_value = mock_result
|
||||
|
||||
repo = ArtifactRepository(mock_db)
|
||||
artifact, created = repo.get_or_create_artifact(
|
||||
sha256="abc123def456",
|
||||
size=1024,
|
||||
filename="test.whl",
|
||||
)
|
||||
|
||||
assert created is False
|
||||
|
||||
def test_get_cached_url_with_artifact_returns_tuple(self):
|
||||
"""get_cached_url_with_artifact should return (CachedUrl, Artifact) tuple."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_cached_url = MagicMock()
|
||||
mock_artifact = MagicMock()
|
||||
mock_db.query.return_value.join.return_value.filter.return_value.first.return_value = (
|
||||
mock_cached_url,
|
||||
mock_artifact,
|
||||
)
|
||||
|
||||
repo = ArtifactRepository(mock_db)
|
||||
result = repo.get_cached_url_with_artifact("url_hash_123")
|
||||
|
||||
assert result == (mock_cached_url, mock_artifact)
|
||||
|
||||
def test_get_cached_url_with_artifact_returns_none_when_not_found(self):
|
||||
"""get_cached_url_with_artifact should return None when URL not cached."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.join.return_value.filter.return_value.first.return_value = None
|
||||
|
||||
repo = ArtifactRepository(mock_db)
|
||||
result = repo.get_cached_url_with_artifact("nonexistent_hash")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_get_artifact_dependencies_returns_list(self):
|
||||
"""get_artifact_dependencies should return list of dependencies."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_dep1 = MagicMock()
|
||||
mock_dep2 = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = [
|
||||
mock_dep1,
|
||||
mock_dep2,
|
||||
]
|
||||
|
||||
repo = ArtifactRepository(mock_db)
|
||||
result = repo.get_artifact_dependencies("artifact_hash_123")
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0] == mock_dep1
|
||||
assert result[1] == mock_dep2
|
||||
|
||||
def test_get_artifact_dependencies_returns_empty_list(self):
|
||||
"""get_artifact_dependencies should return empty list when no dependencies."""
|
||||
from backend.app.db_utils import ArtifactRepository
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = []
|
||||
|
||||
repo = ArtifactRepository(mock_db)
|
||||
result = repo.get_artifact_dependencies("artifact_without_deps")
|
||||
|
||||
assert result == []
|
||||
Reference in New Issue
Block a user