From 6c556104f664e4c1311a5ff6f749da22bede5969 Mon Sep 17 00:00:00 2001 From: Polo Date: Wed, 25 Feb 2026 15:18:40 +0200 Subject: [PATCH 1/2] Add version management endpoints and client methods --- app/api/v1/endpoints/artifacts.py | 78 +++++++++++- app/grpc/clients/artifact_storage_client.py | 70 +++++++++++ tests/unit/test_version_endpoints.py | 125 ++++++++++++++++++++ tests/unit/test_version_management.py | 93 +++++++++++++++ 4 files changed, 365 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_version_endpoints.py create mode 100644 tests/unit/test_version_management.py diff --git a/app/api/v1/endpoints/artifacts.py b/app/api/v1/endpoints/artifacts.py index f449670..347ce40 100644 --- a/app/api/v1/endpoints/artifacts.py +++ b/app/api/v1/endpoints/artifacts.py @@ -173,17 +173,28 @@ async def get_artifact_metadata( async def list_artifacts( project_id: Optional[str] = Query(None, description="Filter by project ID"), pipeline_id: Optional[str] = Query(None, description="Filter by pipeline ID"), + include_versions: bool = Query(False, description="Include version information"), current_user: User = Depends(deps.get_current_user), client: ArtifactStorageClient = Depends(get_artifact_storage_client) ): """ - List artifacts with optional filters. + List artifacts with optional filters and version information. """ try: artifacts = await client.list_artifacts( project_id=project_id, pipeline_id=pipeline_id ) + + # If version information requested, fetch versions for each artifact + if include_versions: + for artifact in artifacts: + try: + artifact["versions"] = await client.get_versions(artifact["id"]) + except Exception: + # If we can't get versions, don't fail the entire request + artifact["versions"] = [] + return artifacts except Exception as e: raise HTTPException( @@ -367,3 +378,68 @@ async def complete_multipart_upload( status_code=500, detail=f"Failed to complete multipart upload: {str(e)}" ) + + +# Version Management Request/Response Models +class ArtifactVersion(BaseModel): + """Artifact version information.""" + version_id: Optional[str] = Field(None, description="S3 version ID") + size: int = Field(..., description="Version size in bytes") + checksum_sha256: Optional[str] = Field(None, description="SHA256 checksum") + created_at: Optional[str] = Field(None, description="Creation timestamp") + is_latest: bool = Field(False, description="Whether this is the latest version") + + +class RestoreArtifactRequest(BaseModel): + """Request to restore an artifact.""" + version_id: Optional[str] = Field(None, description="Target version ID (optional)") + + +class RestoreArtifactResponse(BaseModel): + """Response with restoration status.""" + status: str = Field(..., description="Restoration status") + restored_version_id: Optional[str] = Field(None, description="Restored version ID") + restoration_time: Optional[str] = Field(None, description="Restoration timestamp") + + +# Version Management Endpoints +@router.get("/{artifact_id}/versions", response_model=list[ArtifactVersion], status_code=status.HTTP_200_OK) +async def get_artifact_versions( + artifact_id: str, + current_user: User = Depends(deps.get_current_user), + client: ArtifactStorageClient = Depends(get_artifact_storage_client) +): + """ + Get version history for an artifact. + """ + try: + versions = await client.get_versions(artifact_id) + return versions + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to get artifact versions: {str(e)}" + ) + + +@router.post("/{artifact_id}/restore", response_model=RestoreArtifactResponse, status_code=status.HTTP_200_OK) +async def restore_artifact( + artifact_id: str, + request: RestoreArtifactRequest, + current_user: User = Depends(deps.get_current_user), + client: ArtifactStorageClient = Depends(get_artifact_storage_client) +): + """ + Restore an artifact to a specific version. + """ + try: + result = await client.restore_artifact( + artifact_id=artifact_id, + version_id=request.version_id + ) + return result + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to restore artifact: {str(e)}" + ) diff --git a/app/grpc/clients/artifact_storage_client.py b/app/grpc/clients/artifact_storage_client.py index c0d0b37..8c207f0 100644 --- a/app/grpc/clients/artifact_storage_client.py +++ b/app/grpc/clients/artifact_storage_client.py @@ -453,6 +453,76 @@ async def complete_multipart_upload(self, artifact_id: str, upload_id: str, part logger.error(f"gRPC error completing multipart upload: {e.code()} - {e.details()}") raise + async def get_versions(self, artifact_id: str) -> list: + """ + Get version history for an artifact. + + Args: + artifact_id: Artifact ID to get versions for + + Returns: + List of artifact versions with metadata + """ + if not self.stub: + raise RuntimeError("Client not connected. Call connect() first.") + + try: + from app.grpc.generated import artifact_pb2 + + request = artifact_pb2.GetVersionsRequest(artifact_id=artifact_id) + + metadata = (("x-api-key", self.api_key),) + response = await self.stub.GetVersions(request, metadata=metadata) + + versions = [] + for version in response.versions: + versions.append({ + "version_id": version.version_id if version.HasField("version_id") else None, + "size": version.size, + "checksum_sha256": version.checksum_sha256 if version.HasField("checksum_sha256") else None, + "created_at": version.created_at.ToDatetime().isoformat() if version.HasField("created_at") else None, + "is_latest": version.is_latest if version.HasField("is_latest") else False + }) + + return versions + except grpc.RpcError as e: + logger.error(f"gRPC error getting artifact versions: {e.code()} - {e.details()}") + raise + + async def restore_artifact(self, artifact_id: str, version_id: str = None) -> Dict[str, Any]: + """ + Restore an artifact to a specific version. + + Args: + artifact_id: Artifact ID to restore + version_id: Target version ID (optional) + + Returns: + Dict with restoration status and details + """ + if not self.stub: + raise RuntimeError("Client not connected. Call connect() first.") + + try: + from app.grpc.generated import artifact_pb2 + + request = artifact_pb2.RestoreArtifactRequest( + artifact_id=artifact_id, + version_id=version_id + ) + + metadata = (("x-api-key", self.api_key),) + response = await self.stub.RestoreArtifact(request, metadata=metadata) + + return { + "status": response.status, + "restored_version_id": response.restored_version_id if response.HasField("restored_version_id") else None, + "restoration_time": response.restoration_time.ToDatetime().isoformat() if response.HasField("restoration_time") else None + } + except grpc.RpcError as e: + logger.error(f"gRPC error restoring artifact: {e.code()} - {e.details()}") + raise + async def close(self): """Close gRPC connection.""" if self.channel: diff --git a/tests/unit/test_version_endpoints.py b/tests/unit/test_version_endpoints.py new file mode 100644 index 0000000..b0aa599 --- /dev/null +++ b/tests/unit/test_version_endpoints.py @@ -0,0 +1,125 @@ +"""Unit tests for artifact version management endpoints.""" + +import pytest +from unittest.mock import AsyncMock, patch +from app.api.v1.endpoints.artifacts import ( + get_artifact_versions, + restore_artifact, + ArtifactVersion, + RestoreArtifactRequest, + RestoreArtifactResponse +) + + +@pytest.mark.asyncio +async def test_get_artifact_versions_endpoint_success(): + """Test successful artifact versions endpoint.""" + mock_client = AsyncMock() + mock_client.get_versions.return_value = [ + { + "version_id": "version-123", + "size": 1024, + "checksum_sha256": "abc123", + "created_at": "2026-02-25T16:00:00", + "is_latest": True + } + ] + + with patch("app.api.v1.endpoints.artifacts.get_artifact_storage_client", return_value=mock_client), \ + patch("app.api.v1.endpoints.artifacts.deps.get_current_user", return_value={"id": 1}): + + result = await get_artifact_versions("test-artifact-id", mock_client, mock_client) + + assert len(result) == 1 + assert result[0]["version_id"] == "version-123" + assert result[0]["is_latest"] is True + mock_client.get_versions.assert_called_once_with("test-artifact-id") + + +@pytest.mark.asyncio +async def test_restore_artifact_endpoint_success(): + """Test successful artifact restoration endpoint.""" + mock_client = AsyncMock() + mock_client.restore_artifact.return_value = { + "status": "restored", + "restored_version_id": "version-123", + "restoration_time": "2026-02-25T16:00:00" + } + + with patch("app.api.v1.endpoints.artifacts.get_artifact_storage_client", return_value=mock_client), \ + patch("app.api.v1.endpoints.artifacts.deps.get_current_user", return_value={"id": 1}): + + request = RestoreArtifactRequest(version_id="version-123") + result = await restore_artifact("test-artifact-id", request, mock_client, mock_client) + + assert result["status"] == "restored" + assert result["restored_version_id"] == "version-123" + mock_client.restore_artifact.assert_called_once_with("test-artifact-id", "version-123") + + +@pytest.mark.asyncio +async def test_restore_artifact_endpoint_latest(): + """Test artifact restoration to latest version.""" + mock_client = AsyncMock() + mock_client.restore_artifact.return_value = { + "status": "restored", + "restored_version_id": None, + "restoration_time": "2026-02-25T16:00:00" + } + + with patch("app.api.v1.endpoints.artifacts.get_artifact_storage_client", return_value=mock_client), \ + patch("app.api.v1.endpoints.artifacts.deps.get_current_user", return_value={"id": 1}): + + request = RestoreArtifactRequest() # No version_id specified + result = await restore_artifact("test-artifact-id", request, mock_client, mock_client) + + assert result["status"] == "restored" + assert result["restored_version_id"] is None + mock_client.restore_artifact.assert_called_once_with("test-artifact-id", None) + + +@pytest.mark.asyncio +async def test_version_management_endpoint_error_handling(): + """Test version management endpoint error handling.""" + mock_client = AsyncMock() + mock_client.get_versions.side_effect = Exception("Storage error") + + with patch("app.api.v1.endpoints.artifacts.get_artifact_storage_client", return_value=mock_client), \ + patch("app.api.v1.endpoints.artifacts.deps.get_current_user", return_value={"id": 1}): + + with pytest.raises(Exception): # Should raise HTTPException + await get_artifact_versions("test-artifact-id", mock_client, mock_client) + + +# Test request/response models +def test_artifact_version_model(): + """Test ArtifactVersion model validation.""" + version = ArtifactVersion( + version_id="version-123", + size=1024, + checksum_sha256="abc123", + created_at="2026-02-25T16:00:00", + is_latest=True + ) + assert version.version_id == "version-123" + assert version.size == 1024 + assert version.checksum_sha256 == "abc123" + assert version.is_latest is True + + +def test_restore_artifact_request_model(): + """Test RestoreArtifactRequest model validation.""" + request = RestoreArtifactRequest(version_id="version-123") + assert request.version_id == "version-123" + + +def test_restore_artifact_response_model(): + """Test RestoreArtifactResponse model validation.""" + response = RestoreArtifactResponse( + status="restored", + restored_version_id="version-123", + restoration_time="2026-02-25T16:00:00" + ) + assert response.status == "restored" + assert response.restored_version_id == "version-123" + assert response.restoration_time == "2026-02-25T16:00:00" diff --git a/tests/unit/test_version_management.py b/tests/unit/test_version_management.py new file mode 100644 index 0000000..014de58 --- /dev/null +++ b/tests/unit/test_version_management.py @@ -0,0 +1,93 @@ +"""Unit tests for artifact version management.""" + +import pytest +from unittest.mock import AsyncMock, patch, MagicMock +from app.grpc.clients.artifact_storage_client import ArtifactStorageClient + + +@pytest.mark.asyncio +async def test_get_versions_success(): + """Test successful artifact version retrieval.""" + client = ArtifactStorageClient("localhost", 50051, "test-key") + + with patch("app.grpc.clients.artifact_storage_client.grpc.aio.insecure_channel"), \ + patch("app.grpc.generated.artifact_pb2_grpc.ArtifactServiceStub") as mock_stub: + + # Mock successful response + mock_response = MagicMock() + mock_version = MagicMock() + mock_version.version_id = "version-123" + mock_version.size = 1024 + mock_version.checksum_sha256 = "abc123" + mock_version.created_at.ToDatetime.return_value.isoformat.return_value = "2026-02-25T16:00:00" + mock_version.is_latest = True + mock_version.HasField.return_value = True + + mock_response.versions = [mock_version] + mock_stub.return_value.GetVersions.return_value = mock_response + + await client.connect() + result = await client.get_versions("test-artifact-id") + + assert len(result) == 1 + assert result[0]["version_id"] == "version-123" + assert result[0]["size"] == 1024 + assert result[0]["checksum_sha256"] == "abc123" + assert result[0]["is_latest"] is True + + +@pytest.mark.asyncio +async def test_restore_artifact_success(): + """Test successful artifact restoration.""" + client = ArtifactStorageClient("localhost", 50051, "test-key") + + with patch("app.grpc.clients.artifact_storage_client.grpc.aio.insecure_channel"), \ + patch("app.grpc.generated.artifact_pb2_grpc.ArtifactServiceStub") as mock_stub: + + # Mock successful response + mock_response = MagicMock() + mock_response.status = "restored" + mock_response.restored_version_id = "version-123" + mock_response.HasField.return_value = True + mock_response.restoration_time.ToDatetime.return_value.isoformat.return_value = "2026-02-25T16:00:00" + mock_stub.return_value.RestoreArtifact.return_value = mock_response + + await client.connect() + result = await client.restore_artifact("test-artifact-id", "version-123") + + assert result["status"] == "restored" + assert result["restored_version_id"] == "version-123" + assert result["restoration_time"] == "2026-02-25T16:00:00" + + +@pytest.mark.asyncio +async def test_restore_artifact_latest(): + """Test artifact restoration to latest version.""" + client = ArtifactStorageClient("localhost", 50051, "test-key") + + with patch("app.grpc.clients.artifact_storage_client.grpc.aio.insecure_channel"), \ + patch("app.grpc.generated.artifact_pb2_grpc.ArtifactServiceStub") as mock_stub: + + # Mock successful response + mock_response = MagicMock() + mock_response.status = "restored" + mock_response.HasField.return_value = False # No version_id specified + mock_stub.return_value.RestoreArtifact.return_value = mock_response + + await client.connect() + result = await client.restore_artifact("test-artifact-id", None) + + assert result["status"] == "restored" + assert result["restored_version_id"] is None + + +@pytest.mark.asyncio +async def test_version_management_error_handling(): + """Test version management error handling.""" + client = ArtifactStorageClient("localhost", 50051, "test-key") + + with pytest.raises(RuntimeError, match="Client not connected"): + await client.get_versions("test-id") + + with pytest.raises(RuntimeError, match="Client not connected"): + await client.restore_artifact("test-id", "version-123") From 1c2b56d3482e0a95f061454f594fe02e7dca9f76 Mon Sep 17 00:00:00 2001 From: Polo Date: Wed, 25 Feb 2026 15:24:24 +0200 Subject: [PATCH 2/2] Add retention policy and storage class management endpoints --- app/api/v1/endpoints/artifacts.py | 85 ++++++++++++++++ app/grpc/clients/artifact_storage_client.py | 93 ++++++++++++++++++ tests/unit/test_advanced_endpoints.py | 102 ++++++++++++++++++++ tests/unit/test_advanced_features.py | 62 ++++++++++++ 4 files changed, 342 insertions(+) create mode 100644 tests/unit/test_advanced_endpoints.py create mode 100644 tests/unit/test_advanced_features.py diff --git a/app/api/v1/endpoints/artifacts.py b/app/api/v1/endpoints/artifacts.py index 347ce40..a96f70e 100644 --- a/app/api/v1/endpoints/artifacts.py +++ b/app/api/v1/endpoints/artifacts.py @@ -443,3 +443,88 @@ async def restore_artifact( status_code=500, detail=f"Failed to restore artifact: {str(e)}" ) + + +# Retention Policy Request/Response Models +class SetRetentionPolicyRequest(BaseModel): + """Request to set project retention policy.""" + project_id: str = Field(..., description="Project ID") + max_bytes: int = Field(..., description="Maximum storage bytes", ge=0) + ttl_days: int = Field(..., description="Default TTL in days", ge=1) + + +class RetentionPolicyResponse(BaseModel): + """Response with retention policy details.""" + project_id: str = Field(..., description="Project ID") + max_bytes: Optional[int] = Field(None, description="Maximum storage bytes") + ttl_days: Optional[int] = Field(None, description="Default TTL in days") + current_usage: Optional[int] = Field(None, description="Current usage in bytes") + + +class StorageClassInfo(BaseModel): + """Storage class information.""" + name: str = Field(..., description="Storage class name") + description: str = Field(..., description="Storage class description") + cost_per_gb: Optional[float] = Field(None, description="Cost per GB") + + +# Retention Policy Endpoints +@router.put("/projects/{project_id}/retention", response_model=RetentionPolicyResponse, status_code=status.HTTP_200_OK) +async def set_retention_policy( + project_id: str, + request: SetRetentionPolicyRequest, + current_user: User = Depends(deps.get_current_user), + client: ArtifactStorageClient = Depends(get_artifact_storage_client) +): + """ + Set retention policy for a project. + """ + try: + result = await client.set_retention_policy( + project_id=project_id, + max_bytes=request.max_bytes, + ttl_days=request.ttl_days + ) + return result + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to set retention policy: {str(e)}" + ) + + +@router.get("/projects/{project_id}/retention", response_model=RetentionPolicyResponse, status_code=status.HTTP_200_OK) +async def get_retention_policy( + project_id: str, + current_user: User = Depends(deps.get_current_user), + client: ArtifactStorageClient = Depends(get_artifact_storage_client) +): + """ + Get retention policy for a project. + """ + try: + result = await client.get_retention_policy(project_id) + return result + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to get retention policy: {str(e)}" + ) + + +@router.get("/storage-classes", response_model=list[StorageClassInfo], status_code=status.HTTP_200_OK) +async def get_storage_classes( + current_user: User = Depends(deps.get_current_user), + client: ArtifactStorageClient = Depends(get_artifact_storage_client) +): + """ + Get available storage classes. + """ + try: + result = await client.get_storage_classes() + return result + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to get storage classes: {str(e)}" + ) diff --git a/app/grpc/clients/artifact_storage_client.py b/app/grpc/clients/artifact_storage_client.py index 8c207f0..25722f8 100644 --- a/app/grpc/clients/artifact_storage_client.py +++ b/app/grpc/clients/artifact_storage_client.py @@ -523,6 +523,99 @@ async def restore_artifact(self, artifact_id: str, version_id: str = None) -> Di logger.error(f"gRPC error restoring artifact: {e.code()} - {e.details()}") raise + async def set_retention_policy(self, project_id: str, max_bytes: int, ttl_days: int) -> Dict[str, Any]: + """ + Set retention policy for a project. + + Args: + project_id: Project ID to set policy for + max_bytes: Maximum storage bytes allowed + ttl_days: Default TTL for artifacts in days + + Returns: + Dict with policy status and details + """ + if not self.stub: + raise RuntimeError("Client not connected. Call connect() first.") + + try: + from app.grpc.generated import artifact_pb2 + + request = artifact_pb2.SetRetentionPolicyRequest( + project_id=project_id, + max_bytes=max_bytes, + ttl_days=ttl_days + ) + + metadata = (("x-api-key", self.api_key),) + response = await self.stub.SetRetentionPolicy(request, metadata=metadata) + + return { + "status": response.status, + "project_id": response.project_id, + "max_bytes": response.max_bytes if response.HasField("max_bytes") else None, + "ttl_days": response.ttl_days if response.HasField("ttl_days") else None + } + except grpc.RpcError as e: + logger.error(f"gRPC error setting retention policy: {e.code()} - {e.details()}") + raise + + async def get_retention_policy(self, project_id: str) -> Dict[str, Any]: + """ + Get retention policy for a project. + + Args: + project_id: Project ID to get policy for + + Returns: + Dict with policy details + """ + if not self.stub: + raise RuntimeError("Client not connected. Call connect() first.") + + try: + from app.grpc.generated import artifact_pb2 + + request = artifact_pb2.GetRetentionPolicyRequest(project_id=project_id) + + metadata = (("x-api-key", self.api_key),) + response = await self.stub.GetRetentionPolicy(request, metadata=metadata) + + return { + "project_id": response.project_id, + "max_bytes": response.max_bytes if response.HasField("max_bytes") else None, + "ttl_days": response.ttl_days if response.HasField("ttl_days") else None, + "current_usage": response.current_usage if response.HasField("current_usage") else None + } + except grpc.RpcError as e: + logger.error(f"gRPC error getting retention policy: {e.code()} - {e.details()}") + raise + + async def get_storage_classes(self) -> Dict[str, Any]: + """ + Get available storage classes. + + Returns: + Dict with available storage classes + """ + if not self.stub: + raise RuntimeError("Client not connected. Call connect() first.") + + try: + from app.grpc.generated import artifact_pb2 + + request = artifact_pb2.GetStorageClassesRequest() + + metadata = (("x-api-key", self.api_key),) + response = await self.stub.GetStorageClasses(request, metadata=metadata) + + return { + "storage_classes": list(response.storage_classes) + } + except grpc.RpcError as e: + logger.error(f"gRPC error getting storage classes: {e.code()} - {e.details()}") + raise + async def close(self): """Close gRPC connection.""" if self.channel: diff --git a/tests/unit/test_advanced_endpoints.py b/tests/unit/test_advanced_endpoints.py new file mode 100644 index 0000000..819f5fa --- /dev/null +++ b/tests/unit/test_advanced_endpoints.py @@ -0,0 +1,102 @@ +"""Unit tests for advanced features endpoints.""" + +import pytest +from unittest.mock import AsyncMock, patch +from app.api.v1.endpoints.artifacts import ( + set_retention_policy, + get_retention_policy, + get_storage_classes, + SetRetentionPolicyRequest, + RetentionPolicyResponse, + StorageClassInfo +) + + +@pytest.mark.asyncio +async def test_set_retention_policy_endpoint_success(): + """Test successful retention policy setting endpoint.""" + mock_client = AsyncMock() + mock_client.set_retention_policy.return_value = { + "status": "set", + "project_id": "test-project", + "max_bytes": 1000000000, + "ttl_days": 30 + } + + with patch("app.api.v1.endpoints.artifacts.get_artifact_storage_client", return_value=mock_client), \ + patch("app.api.v1.endpoints.artifacts.deps.get_current_user", return_value={"id": 1}): + + request = SetRetentionPolicyRequest( + project_id="test-project", + max_bytes=1000000000, + ttl_days=30 + ) + + result = await set_retention_policy("test-project", request, mock_client, mock_client) + + assert result["status"] == "set" + assert result["project_id"] == "test-project" + assert result["max_bytes"] == 1000000000 + assert result["ttl_days"] == 30 + mock_client.set_retention_policy.assert_called_once_with("test-project", 1000000000, 30) + + +@pytest.mark.asyncio +async def test_get_storage_classes_endpoint_success(): + """Test successful storage classes endpoint.""" + mock_client = AsyncMock() + mock_client.get_storage_classes.return_value = { + "storage_classes": ["standard", "premium", "archive"] + } + + with patch("app.api.v1.endpoints.artifacts.get_artifact_storage_client", return_value=mock_client), \ + patch("app.api.v1.endpoints.artifacts.deps.get_current_user", return_value={"id": 1}): + + result = await get_storage_classes(mock_client, mock_client) + + assert result["storage_classes"] == ["standard", "premium", "archive"] + mock_client.get_storage_classes.assert_called_once() + + +@pytest.mark.asyncio +async def test_advanced_features_endpoint_error_handling(): + """Test advanced features endpoint error handling.""" + mock_client = AsyncMock() + mock_client.set_retention_policy.side_effect = Exception("Storage error") + + with patch("app.api.v1.endpoints.artifacts.get_artifact_storage_client", return_value=mock_client), \ + patch("app.api.v1.endpoints.artifacts.deps.get_current_user", return_value={"id": 1}): + + request = SetRetentionPolicyRequest( + project_id="test-project", + max_bytes=1000000000, + ttl_days=30 + ) + + with pytest.raises(Exception): # Should raise HTTPException + await set_retention_policy("test-project", request, mock_client, mock_client) + + +# Test request/response models +def test_set_retention_policy_request_model(): + """Test SetRetentionPolicyRequest model validation.""" + request = SetRetentionPolicyRequest( + project_id="test-project", + max_bytes=1000000000, + ttl_days=30 + ) + assert request.project_id == "test-project" + assert request.max_bytes == 1000000000 + assert request.ttl_days == 30 + + +def test_storage_class_info_model(): + """Test StorageClassInfo model validation.""" + storage_class = StorageClassInfo( + name="standard", + description="Standard storage class", + cost_per_gb=0.023 + ) + assert storage_class.name == "standard" + assert storage_class.description == "Standard storage class" + assert storage_class.cost_per_gb == 0.023 diff --git a/tests/unit/test_advanced_features.py b/tests/unit/test_advanced_features.py new file mode 100644 index 0000000..7de41b1 --- /dev/null +++ b/tests/unit/test_advanced_features.py @@ -0,0 +1,62 @@ +"""Unit tests for advanced features.""" + +import pytest +from unittest.mock import AsyncMock, patch +from app.grpc.clients.artifact_storage_client import ArtifactStorageClient + + +@pytest.mark.asyncio +async def test_set_retention_policy_success(): + """Test successful retention policy setting.""" + client = ArtifactStorageClient("localhost", 50051, "test-key") + + with patch("app.grpc.clients.artifact_storage_client.grpc.aio.insecure_channel"), \ + patch("app.grpc.generated.artifact_pb2_grpc.ArtifactServiceStub") as mock_stub: + + # Mock successful response + mock_response = MagicMock() + mock_response.status = "set" + mock_response.project_id = "test-project" + mock_response.max_bytes = 1000000000 + mock_response.ttl_days = 30 + mock_response.HasField.return_value = True + mock_stub.return_value.SetRetentionPolicy.return_value = mock_response + + await client.connect() + result = await client.set_retention_policy("test-project", 1000000000, 30) + + assert result["status"] == "set" + assert result["project_id"] == "test-project" + assert result["max_bytes"] == 1000000000 + assert result["ttl_days"] == 30 + + +@pytest.mark.asyncio +async def test_get_storage_classes_success(): + """Test successful storage classes retrieval.""" + client = ArtifactStorageClient("localhost", 50051, "test-key") + + with patch("app.grpc.clients.artifact_storage_client.grpc.aio.insecure_channel"), \ + patch("app.grpc.generated.artifact_pb2_grpc.ArtifactServiceStub") as mock_stub: + + # Mock successful response + mock_response = MagicMock() + mock_response.storage_classes = ["standard", "premium", "archive"] + mock_stub.return_value.GetStorageClasses.return_value = mock_response + + await client.connect() + result = await client.get_storage_classes() + + assert result["storage_classes"] == ["standard", "premium", "archive"] + + +@pytest.mark.asyncio +async def test_advanced_features_error_handling(): + """Test advanced features error handling.""" + client = ArtifactStorageClient("localhost", 50051, "test-key") + + with pytest.raises(RuntimeError, match="Client not connected"): + await client.set_retention_policy("test-project", 1000000000, 30) + + with pytest.raises(RuntimeError, match="Client not connected"): + await client.get_storage_classes()