"""Tests for storage abstraction layer.""" import io from unittest.mock import AsyncMock, MagicMock, patch import pytest from botocore.exceptions import ClientError from reflector.storage.base import StoragePermissionError from reflector.storage.storage_aws import AwsStorage @pytest.mark.asyncio async def test_aws_storage_stream_to_fileobj(): """Test that AWS storage can stream directly to a file object without loading into memory.""" # Setup storage = AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) # Mock download_fileobj to write data async def mock_download(Bucket, Key, Fileobj, **kwargs): Fileobj.write(b"chunk1chunk2") mock_client = AsyncMock() mock_client.download_fileobj = AsyncMock(side_effect=mock_download) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) # Patch the session client with patch.object(storage.session, "client", return_value=mock_client): # Create a file-like object to stream to output = io.BytesIO() # Act - stream to file object await storage.stream_to_fileobj("test-file.mp4", output, bucket="test-bucket") # Assert mock_client.download_fileobj.assert_called_once_with( Bucket="test-bucket", Key="test-file.mp4", Fileobj=output ) # Check that data was written to output output.seek(0) assert output.read() == b"chunk1chunk2" @pytest.mark.asyncio async def test_aws_storage_stream_to_fileobj_with_folder(): """Test streaming with folder prefix in bucket name.""" storage = AwsStorage( aws_bucket_name="test-bucket/recordings", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) async def mock_download(Bucket, Key, Fileobj, **kwargs): Fileobj.write(b"data") mock_client = AsyncMock() mock_client.download_fileobj = AsyncMock(side_effect=mock_download) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): output = io.BytesIO() await storage.stream_to_fileobj("file.mp4", output, bucket="other-bucket") # Should use folder prefix from instance config mock_client.download_fileobj.assert_called_once_with( Bucket="other-bucket", Key="recordings/file.mp4", Fileobj=output ) @pytest.mark.asyncio async def test_storage_base_class_stream_to_fileobj(): """Test that base Storage class has stream_to_fileobj method.""" from reflector.storage.base import Storage # Verify method exists in base class assert hasattr(Storage, "stream_to_fileobj") # Create a mock storage instance storage = MagicMock(spec=Storage) storage.stream_to_fileobj = AsyncMock() # Should be callable await storage.stream_to_fileobj("file.mp4", io.BytesIO()) storage.stream_to_fileobj.assert_called_once() @pytest.mark.asyncio async def test_aws_storage_stream_uses_download_fileobj(): """Test that download_fileobj is called correctly.""" storage = AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) async def mock_download(Bucket, Key, Fileobj, **kwargs): Fileobj.write(b"data") mock_client = AsyncMock() mock_client.download_fileobj = AsyncMock(side_effect=mock_download) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): output = io.BytesIO() await storage.stream_to_fileobj("test.mp4", output) # Verify download_fileobj was called with correct parameters mock_client.download_fileobj.assert_called_once_with( Bucket="test-bucket", Key="test.mp4", Fileobj=output ) @pytest.mark.asyncio async def test_aws_storage_handles_access_denied_error(): """Test that AccessDenied errors are caught and wrapped in StoragePermissionError.""" storage = AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) # Mock ClientError with AccessDenied error_response = {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}} mock_client = AsyncMock() mock_client.put_object = AsyncMock( side_effect=ClientError(error_response, "PutObject") ) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): with pytest.raises(StoragePermissionError) as exc_info: await storage.put_file("test.txt", b"data") # Verify error message contains expected information error_msg = str(exc_info.value) assert "AccessDenied" in error_msg assert "default bucket 'test-bucket'" in error_msg assert "S3 upload failed" in error_msg @pytest.mark.asyncio async def test_aws_storage_handles_no_such_bucket_error(): """Test that NoSuchBucket errors are caught and wrapped in StoragePermissionError.""" storage = AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) # Mock ClientError with NoSuchBucket error_response = { "Error": { "Code": "NoSuchBucket", "Message": "The specified bucket does not exist", } } mock_client = AsyncMock() mock_client.delete_object = AsyncMock( side_effect=ClientError(error_response, "DeleteObject") ) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): with pytest.raises(StoragePermissionError) as exc_info: await storage.delete_file("test.txt") # Verify error message contains expected information error_msg = str(exc_info.value) assert "NoSuchBucket" in error_msg assert "default bucket 'test-bucket'" in error_msg assert "S3 delete failed" in error_msg @pytest.mark.asyncio async def test_aws_storage_error_message_with_bucket_override(): """Test that error messages correctly show overridden bucket.""" storage = AwsStorage( aws_bucket_name="default-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) # Mock ClientError with AccessDenied error_response = {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}} mock_client = AsyncMock() mock_client.get_object = AsyncMock( side_effect=ClientError(error_response, "GetObject") ) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): with pytest.raises(StoragePermissionError) as exc_info: await storage.get_file("test.txt", bucket="override-bucket") # Verify error message shows overridden bucket, not default error_msg = str(exc_info.value) assert "overridden bucket 'override-bucket'" in error_msg assert "default-bucket" not in error_msg assert "S3 download failed" in error_msg @pytest.mark.asyncio async def test_aws_storage_reraises_non_handled_errors(): """Test that non-AccessDenied/NoSuchBucket errors are re-raised as-is.""" storage = AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) # Mock ClientError with different error code error_response = { "Error": {"Code": "InternalError", "Message": "Internal Server Error"} } mock_client = AsyncMock() mock_client.put_object = AsyncMock( side_effect=ClientError(error_response, "PutObject") ) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): # Should raise ClientError, not StoragePermissionError with pytest.raises(ClientError) as exc_info: await storage.put_file("test.txt", b"data") # Verify it's the original ClientError assert exc_info.value.response["Error"]["Code"] == "InternalError" @pytest.mark.asyncio async def test_aws_storage_presign_url_handles_errors(): """Test that presigned URL generation handles permission errors.""" storage = AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) # Mock ClientError with AccessDenied during presign operation error_response = {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}} mock_client = AsyncMock() mock_client.generate_presigned_url = AsyncMock( side_effect=ClientError(error_response, "GeneratePresignedUrl") ) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): with pytest.raises(StoragePermissionError) as exc_info: await storage.get_file_url("test.txt") # Verify error message error_msg = str(exc_info.value) assert "S3 presign failed" in error_msg assert "AccessDenied" in error_msg @pytest.mark.asyncio async def test_aws_storage_list_objects_handles_errors(): """Test that list_objects handles permission errors.""" storage = AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", ) # Mock ClientError during list operation error_response = {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}} mock_paginator = MagicMock() async def mock_paginate(*args, **kwargs): raise ClientError(error_response, "ListObjectsV2") yield # Make it an async generator mock_paginator.paginate = mock_paginate mock_client = AsyncMock() mock_client.get_paginator = MagicMock(return_value=mock_paginator) mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) with patch.object(storage.session, "client", return_value=mock_client): with pytest.raises(StoragePermissionError) as exc_info: await storage.list_objects(prefix="test/") error_msg = str(exc_info.value) assert "S3 list_objects failed" in error_msg assert "AccessDenied" in error_msg def test_aws_storage_constructor_rejects_mixed_auth(): """Test that constructor rejects both role_arn and access keys.""" with pytest.raises(ValueError, match="cannot use both.*role_arn.*access keys"): AwsStorage( aws_bucket_name="test-bucket", aws_region="us-east-1", aws_access_key_id="test-key", aws_secret_access_key="test-secret", aws_role_arn="arn:aws:iam::123456789012:role/test-role", )