chainlit--chainlit/backend/tests/data/storage_clients/test_gcs.py

from unittest.mock import MagicMock, patch

import pytest

from chainlit.data.storage_clients.base import storage_expiry_time
from chainlit.data.storage_clients.gcs import GCSStorageClient


@pytest.fixture
def mock_gcs_client():
    """Create a mock Google Cloud Storage client."""
    # First mock the service_account
    with patch(
        "chainlit.data.storage_clients.gcs.service_account"
    ) as mock_service_account:
        mock_credentials = MagicMock()
        mock_service_account.Credentials.from_service_account_info.return_value = (
            mock_credentials
        )

        # Then mock the storage client
        with patch("chainlit.data.storage_clients.gcs.storage") as mock_storage:
            mock_client = MagicMock()
            mock_storage.Client.return_value = mock_client
            mock_bucket = MagicMock()
            mock_client.bucket.return_value = mock_bucket
            mock_blob = MagicMock()
            mock_bucket.blob.return_value = mock_blob

            yield {
                "storage": mock_storage,
                "client": mock_client,
                "bucket": mock_bucket,
                "blob": mock_blob,
                "service_account": mock_service_account,
                "credentials": mock_credentials,
            }


class TestGCSStorageClient:
    def test_init(self, mock_gcs_client):
        """Test initialization of GCS client."""
        # Remove client assignment, directly call the constructor
        GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Verify service account credentials were created correctly
        mock_gcs_client[
            "service_account"
        ].Credentials.from_service_account_info.assert_called_once_with(
            {
                "type": "service_account",
                "project_id": "test-project",
                "private_key": "test-key",
                "client_email": "test@example.com",
                "token_uri": "https://oauth2.googleapis.com/token",
            }
        )

        # Verify storage client was initialized with credentials
        mock_gcs_client["storage"].Client.assert_called_once_with(
            project="test-project", credentials=mock_gcs_client["credentials"]
        )

        # Verify bucket was retrieved
        mock_gcs_client["client"].bucket.assert_called_once_with("test-bucket")

    def test_sync_get_read_url(self, mock_gcs_client):
        """Test generating a signed URL."""
        # Set up the mock to return a URL
        mock_gcs_client[
            "blob"
        ].generate_signed_url.return_value = "https://signed-url.example.com"

        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset mocks to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Test the method
        url = client.sync_get_read_url("test/path/file.txt")

        # Verify correct methods were called
        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
        mock_gcs_client["blob"].generate_signed_url.assert_called_once_with(
            version="v4", expiration=storage_expiry_time, method="GET"
        )

        assert url == "https://signed-url.example.com"

    @pytest.mark.asyncio
    async def test_get_read_url(self, mock_gcs_client):
        """Test the async wrapper for getting a read URL."""
        # Set up the mock to return a URL
        mock_gcs_client[
            "blob"
        ].generate_signed_url.return_value = "https://signed-url.example.com"

        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset mocks to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Test the async method
        url = await client.get_read_url("test/path/file.txt")

        # Verify correct methods were called
        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
        mock_gcs_client["blob"].generate_signed_url.assert_called_once_with(
            version="v4", expiration=storage_expiry_time, method="GET"
        )

        assert url == "https://signed-url.example.com"

    def test_sync_upload_file(self, mock_gcs_client):
        """Test uploading a file to GCS."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Mock the bucket name property
        mock_gcs_client["bucket"].name = "test-bucket"

        # Test with binary data
        binary_data = b"test content"
        object_key = "test/path/file.txt"

        # Remove the unused result assignment
        client.sync_upload_file(
            object_key=object_key, data=binary_data, mime="text/plain", overwrite=True
        )

        # Check that the blob was called with the correct object key (using assert_any_call instead of assert_called_once_with)
        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
            binary_data, content_type="text/plain"
        )

    def test_sync_upload_file_string_data(self, mock_gcs_client):
        """Test uploading string data to GCS."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()
        mock_gcs_client["bucket"].name = "test-bucket"

        # Test with string data that should be encoded
        string_data = "test content"
        object_key = "test/path/file.txt"

        # Remove the unused result assignment
        client.sync_upload_file(
            object_key=object_key, data=string_data, mime="text/plain", overwrite=True
        )

        # Check that the correct methods were called
        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
            b"test content", content_type="text/plain"
        )

    def test_sync_upload_file_no_overwrite(self, mock_gcs_client):
        """Test upload with overwrite=False and existing file."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Configure blob to indicate file exists
        mock_gcs_client["blob"].exists.return_value = True

        with pytest.raises(
            Exception,
            match=r"Failed to upload file to GCS: File test/path/existing\.txt already exists and overwrite is False",
        ):
            client.sync_upload_file(
                object_key="test/path/existing.txt",
                data=b"test content",
                overwrite=False,
            )

        mock_gcs_client["blob"].exists.assert_called_once()
        mock_gcs_client["blob"].upload_from_string.assert_not_called()

    def test_sync_upload_file_error(self, mock_gcs_client):
        """Test error handling during upload."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Configure upload to throw an exception
        mock_gcs_client["blob"].upload_from_string.side_effect = ValueError(
            "Upload failed"
        )

        with pytest.raises(
            Exception, match="Failed to upload file to GCS: Upload failed"
        ):
            client.sync_upload_file(
                object_key="test/path/file.txt", data=b"test content"
            )

    @pytest.mark.asyncio
    async def test_upload_file(self, mock_gcs_client):
        """Test the async wrapper for uploading a file."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()
        mock_gcs_client["bucket"].name = "test-bucket"

        # Test with binary data
        binary_data = b"test content"
        object_key = "test/path/file.txt"

        # Remove the unused result assignment
        await client.upload_file(
            object_key=object_key, data=binary_data, mime="text/plain", overwrite=True
        )

        # Check that the correct methods were called
        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
            binary_data, content_type="text/plain"
        )

    def test_sync_delete_file(self, mock_gcs_client):
        """Test deleting a file from GCS."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Test successful delete
        result = client.sync_delete_file("test/path/file.txt")

        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
        mock_gcs_client["blob"].delete.assert_called_once()
        assert result is True

    def test_sync_delete_file_error(self, mock_gcs_client):
        """Test error handling during file deletion."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Configure delete to throw an exception
        mock_gcs_client["blob"].delete.side_effect = ValueError("Delete failed")

        # The method should catch the exception and return False
        result = client.sync_delete_file("test/path/file.txt")

        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
        mock_gcs_client["blob"].delete.assert_called_once()
        assert result is False

    @pytest.mark.asyncio
    async def test_delete_file(self, mock_gcs_client):
        """Test the async wrapper for deleting a file."""
        client = GCSStorageClient(
            bucket_name="test-bucket",
            project_id="test-project",
            client_email="test@example.com",
            private_key="test-key",
        )

        # Reset the mock to ensure clean state
        mock_gcs_client["bucket"].reset_mock()
        mock_gcs_client["blob"].reset_mock()

        # Test successful delete
        result = await client.delete_file("test/path/file.txt")

        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
        mock_gcs_client["blob"].delete.assert_called_once()
        assert result is True