"""Tests for abacus_file_upload.logic.file_upload_helpers module.""" from datetime import datetime, timezone from unittest.mock import Mock, patch import pytest from botocore.exceptions import ClientError from abacus_file_upload.constants import S3_MAX_PARTS, UPLOAD_STATUSES from abacus_file_upload.logic.file_upload_helpers import ( cleanup_s3_upload, handle_multipart_upload, handle_single_part_upload, handle_upload_operation_error, quarantine_s3_upload, validate_with_upload_config, verify_s3_upload, ) from abacus_file_upload.models import FileUpload, FileUploadConfig from core.config import Config class TestCleanupS3Upload: """Tests for cleanup_s3_upload function.""" @pytest.fixture def mock_s3_connector(self): """Create a mock S3 connector.""" return Mock() @pytest.fixture def file_upload_multipart(self): """Create a FileUpload with multipart upload.""" upload = Mock(spec=FileUpload) upload.multipart_upload_id = 'test-upload-id' upload.s3_bucket = 'test-bucket' upload.s3_key = 'test-key' return upload @pytest.fixture def file_upload_single(self): """Create a FileUpload without multipart upload.""" upload = Mock(spec=FileUpload) upload.multipart_upload_id = None upload.s3_bucket = 'test-bucket' upload.s3_key = 'test-key' return upload def test_cleanup_multipart_upload_success( self, mock_s3_connector, file_upload_multipart ): """Test cleanup successfully aborts multipart upload.""" cleanup_s3_upload(mock_s3_connector, file_upload_multipart) mock_s3_connector.abort_multipart_upload.assert_called_once_with( 'test-bucket', 'test-key', 'test-upload-id' ) def test_cleanup_multipart_upload_failure( self, mock_s3_connector, file_upload_multipart ): """Test cleanup handles abort failure gracefully.""" mock_s3_connector.abort_multipart_upload.side_effect = ClientError( {'Error': {'Code': '500', 'Message': 'Internal Error'}}, 'abort_multipart_upload', ) # Should not raise exception cleanup_s3_upload(mock_s3_connector, file_upload_multipart) mock_s3_connector.abort_multipart_upload.assert_called_once() def test_cleanup_single_upload_exists(self, mock_s3_connector, file_upload_single): """Test cleanup deletes single-part upload when file exists.""" mock_s3_connector.object_exists.return_value = True cleanup_s3_upload(mock_s3_connector, file_upload_single) mock_s3_connector.object_exists.assert_called_once_with( 'test-bucket', 'test-key' ) mock_s3_connector.delete_object.assert_called_once_with( 'test-bucket', 'test-key' ) def test_cleanup_single_upload_not_exists( self, mock_s3_connector, file_upload_single ): """Test cleanup skips deletion when file doesn't exist.""" mock_s3_connector.object_exists.return_value = False cleanup_s3_upload(mock_s3_connector, file_upload_single) mock_s3_connector.object_exists.assert_called_once() mock_s3_connector.delete_object.assert_not_called() def test_cleanup_single_upload_delete_failure( self, mock_s3_connector, file_upload_single ): """Test cleanup handles delete failure gracefully.""" mock_s3_connector.object_exists.return_value = True mock_s3_connector.delete_object.side_effect = ClientError( {'Error': {'Code': '500', 'Message': 'Internal Error'}}, 'delete_object' ) # Should not raise exception cleanup_s3_upload(mock_s3_connector, file_upload_single) mock_s3_connector.delete_object.assert_called_once() class TestHandleMultipartUpload: """Tests for handle_multipart_upload function.""" @pytest.fixture def mock_s3_connector(self): """Create a mock S3 connector.""" connector = Mock() connector.initiate_multipart_upload.return_value = 'test-upload-id' connector.generate_multipart_presigned_urls.return_value = [ { 'part_number': 1, 'url': 'https://s3.amazonaws.com/part1', 'expires_at': datetime.now(timezone.utc), }, { 'part_number': 2, 'url': 'https://s3.amazonaws.com/part2', 'expires_at': datetime.now(timezone.utc), }, ] connector.generate_complete_multipart_presigned_url.return_value = ( 'https://s3.amazonaws.com/complete' ) return connector def test_handle_multipart_upload_success(self, mock_s3_connector): """Test successful multipart upload setup.""" result = handle_multipart_upload( mock_s3_connector, 's3-bucket', 's3-key', 100 * 1024 * 1024, # 100MB 10 * 1024 * 1024, # 10MB min chunk {'filename': 'test.csv'}, 'text/csv', ) assert 'parts' in result assert 'complete_url' in result assert 'chunk_size_bytes' in result assert 'multipart_upload_id' in result assert 'total_parts' in result assert len(result['parts']) == 2 assert result['multipart_upload_id'] == 'test-upload-id' assert result['complete_url'] == 'https://s3.amazonaws.com/complete' # Verify S3 calls mock_s3_connector.initiate_multipart_upload.assert_called_once() mock_s3_connector.generate_multipart_presigned_urls.assert_called_once() mock_s3_connector.generate_complete_multipart_presigned_url.assert_called_once() def test_handle_multipart_upload_parts_format(self, mock_s3_connector): """Test multipart upload returns correctly formatted parts.""" result = handle_multipart_upload( mock_s3_connector, 's3-bucket', 's3-key', 100 * 1024 * 1024, 10 * 1024 * 1024, {}, None, ) for part in result['parts']: assert 'part_number' in part assert 'url' in part assert 'expires_at' in part assert isinstance(part['expires_at'], str) def test_handle_multipart_upload_too_many_parts(self, mock_s3_connector): """Test multipart upload raises error for too many parts.""" # File size that would require more than S3_MAX_PARTS file_size = (S3_MAX_PARTS + 1) * 10 * 1024 * 1024 # Exceeds max parts with pytest.raises(ValueError, match='File too large'): handle_multipart_upload( mock_s3_connector, 's3-bucket', 's3-key', file_size, 10 * 1024 * 1024, {}, None, ) def test_handle_multipart_upload_with_metadata(self, mock_s3_connector): """Test multipart upload passes metadata to S3.""" metadata = {'filename': 'test.csv', 'uploadtype': 'adjustments'} handle_multipart_upload( mock_s3_connector, 's3-bucket', 's3-key', 100 * 1024 * 1024, 10 * 1024 * 1024, metadata, 'text/csv', ) call_args = mock_s3_connector.initiate_multipart_upload.call_args assert call_args[1]['metadata'] == metadata assert call_args[1]['content_type'] == 'text/csv' class TestHandleSinglePartUpload: """Tests for handle_single_part_upload function.""" @pytest.fixture def mock_s3_connector(self): """Create a mock S3 connector.""" connector = Mock() connector.generate_put_presigned_url.return_value = ( 'https://s3.amazonaws.com/upload' ) return connector def test_handle_single_part_upload_success(self, mock_s3_connector): """Test successful single-part upload setup.""" result = handle_single_part_upload( mock_s3_connector, 's3-bucket', 's3-key', {'filename': 'test.csv'}, 'abc123def456abc123def456abc123de', 'text/csv', ) assert 'upload_url' in result assert 'required_headers' in result assert result['upload_url'] == 'https://s3.amazonaws.com/upload' headers = result['required_headers'] assert 'Content-MD5' in headers assert 'Content-Type' in headers assert headers['Content-Type'] == 'text/csv' def test_handle_single_part_upload_metadata_headers(self, mock_s3_connector): """Test single-part upload includes metadata as headers.""" metadata = {'filename': 'test.csv', 'uploadtype': 'adjustments'} result = handle_single_part_upload( mock_s3_connector, 's3-bucket', 's3-key', metadata, 'abc123def456abc123def456abc123de', 'text/csv', ) headers = result['required_headers'] assert 'x-amz-meta-filename' in headers assert 'x-amz-meta-uploadtype' in headers assert headers['x-amz-meta-filename'] == 'test.csv' assert headers['x-amz-meta-uploadtype'] == 'adjustments' def test_handle_single_part_upload_no_mime_type(self, mock_s3_connector): """Test single-part upload without MIME type.""" result = handle_single_part_upload( mock_s3_connector, 's3-bucket', 's3-key', {}, 'abc123def456abc123def456abc123de', None, ) headers = result['required_headers'] assert 'Content-MD5' in headers assert 'Content-Type' not in headers class TestValidateWithUploadConfig: """Tests for validate_with_upload_config function.""" @pytest.fixture def config(self): """Create a mock FileUploadConfig.""" cfg = Mock(spec=FileUploadConfig) cfg.allowed_file_types = ['csv', 'xlsx', 'pdf'] cfg.max_file_size_bytes = 10 * 1024 * 1024 # 10MB return cfg def test_validate_success(self, config): """Test successful validation.""" # Should not raise exception validate_with_upload_config(config, 'report.csv', 5 * 1024 * 1024) def test_validate_invalid_file_type(self, config): """Test validation fails for invalid file type.""" with pytest.raises(ValueError, match='File type not allowed'): validate_with_upload_config(config, 'report.txt', 5 * 1024 * 1024) def test_validate_file_too_large(self, config): """Test validation fails for file too large.""" with pytest.raises(ValueError, match='File size exceeds maximum'): validate_with_upload_config(config, 'report.csv', 20 * 1024 * 1024) def test_validate_all_types_allowed(self, config): """Test validation passes when all types allowed.""" config.allowed_file_types = None # Should not raise exception validate_with_upload_config(config, 'report.txt', 5 * 1024 * 1024) def test_validate_uppercase_extension(self, config): """Test validation handles uppercase extensions.""" # Should not raise exception (extension is normalized to lowercase) validate_with_upload_config(config, 'report.CSV', 5 * 1024 * 1024) class TestVerifyS3Upload: """Tests for verify_s3_upload function.""" @pytest.fixture def mock_s3_connector(self): """Create a mock S3 connector.""" connector = Mock() connector.object_exists.return_value = True connector.get_object_metadata.return_value = { 'size': 1000, 'etag': 'abc123', } return connector @pytest.fixture def file_upload_single(self): """Create a FileUpload for single-part upload.""" upload = Mock(spec=FileUpload) upload.s3_bucket = 'test-bucket' upload.s3_key = 'test-key' upload.file_size_bytes = 1000 upload.md5sum = 'abc123' upload.multipart_upload_id = None return upload @pytest.fixture def file_upload_multipart(self): """Create a FileUpload for multipart upload.""" upload = Mock(spec=FileUpload) upload.s3_bucket = 'test-bucket' upload.s3_key = 'test-key' upload.file_size_bytes = 1000 upload.multipart_upload_id = 'test-upload-id' return upload def test_verify_success_single_part(self, mock_s3_connector, file_upload_single): """Test successful verification of single-part upload.""" # Should not raise exception verify_s3_upload(mock_s3_connector, file_upload_single) mock_s3_connector.object_exists.assert_called_once() mock_s3_connector.get_object_metadata.assert_called_once() def test_verify_success_multipart(self, mock_s3_connector, file_upload_multipart): """Test successful verification of multipart upload.""" # Should not raise exception (no MD5 check for multipart) verify_s3_upload(mock_s3_connector, file_upload_multipart) mock_s3_connector.object_exists.assert_called_once() mock_s3_connector.get_object_metadata.assert_called_once() def test_verify_file_not_found(self, mock_s3_connector, file_upload_single): """Test verification fails when file not found.""" mock_s3_connector.object_exists.return_value = False with pytest.raises(ValueError, match='File not found in S3'): verify_s3_upload(mock_s3_connector, file_upload_single) def test_verify_size_mismatch(self, mock_s3_connector, file_upload_single): """Test verification fails when size doesn't match.""" mock_s3_connector.get_object_metadata.return_value = { 'size': 2000, # Different size 'etag': 'abc123', } with pytest.raises(ValueError, match='File size mismatch'): verify_s3_upload(mock_s3_connector, file_upload_single) def test_verify_md5_mismatch(self, mock_s3_connector, file_upload_single): """Test verification fails when MD5 doesn't match.""" mock_s3_connector.get_object_metadata.return_value = { 'size': 1000, 'etag': 'different123', # Different MD5 } with pytest.raises(ValueError, match='MD5 verification failed'): verify_s3_upload(mock_s3_connector, file_upload_single) def test_verify_md5_case_insensitive(self, mock_s3_connector, file_upload_single): """Test MD5 verification is case-insensitive.""" file_upload_single.md5sum = 'ABC123' mock_s3_connector.get_object_metadata.return_value = { 'size': 1000, 'etag': '"abc123"', # With quotes and different case } # Should not raise exception verify_s3_upload(mock_s3_connector, file_upload_single) class TestHandleUploadOperationError: """Tests for handle_upload_operation_error function.""" @pytest.fixture def file_upload(self): """Create a FileUpload mock.""" upload = Mock(spec=FileUpload) upload.update_attributes = Mock() return upload @patch('abacus_file_upload.logic.file_upload_helpers.FileUpload.commit_changes') def test_handle_error_updates_status(self, mock_commit, file_upload): """Test error handler updates upload status to ERROR.""" error = Exception('Test error message') response = handle_upload_operation_error( file_upload, 'test-file-key', error, 'complete', 500 ) # Verify upload was marked as ERROR file_upload.update_attributes.assert_called_once_with( upload_status=UPLOAD_STATUSES.ERROR, error_message='Test error message', ) mock_commit.assert_called_once() # Verify response assert response.status == 500 assert 'Failed to complete upload' in response.message @patch('abacus_file_upload.logic.file_upload_helpers.FileUpload.commit_changes') def test_handle_error_with_different_operations(self, mock_commit, file_upload): """Test error handler works with different operation names.""" error = Exception('Cancel failed') response = handle_upload_operation_error( file_upload, 'test-file-key', error, 'cancel', 500 ) assert response.status == 500 assert 'Failed to cancel upload' in response.message @patch('abacus_file_upload.logic.file_upload_helpers.FileUpload.commit_changes') def test_handle_error_with_custom_status_code(self, mock_commit, file_upload): """Test error handler uses custom status code.""" error = ValueError('Validation failed') response = handle_upload_operation_error( file_upload, 'test-file-key', error, 'validate', 400 ) assert response.status == 400 assert 'Failed to validate upload' in response.message assert 'Validation failed' in response.message @patch('abacus_file_upload.logic.file_upload_helpers.FileUpload.commit_changes') def test_handle_error_when_db_update_fails(self, mock_commit, file_upload): """Test error handler continues when DB update fails.""" file_upload.update_attributes.side_effect = Exception('DB error') error = Exception('Original error') # Should not raise exception even if DB update fails response = handle_upload_operation_error( file_upload, 'test-file-key', error, 'complete', 500 ) # Still returns error response for original error assert response.status == 500 assert 'Failed to complete upload' in response.message assert 'Original error' in response.message class TestQuarantineS3Upload: """Tests for quarantine_s3_upload function.""" @pytest.fixture def mock_s3_connector(self): """Create a mock S3 connector.""" return Mock() @pytest.fixture def file_upload(self): """Create a FileUpload mock.""" upload = Mock(spec=FileUpload) upload.s3_bucket = 'test-bucket' upload.s3_key = 'test-key' return upload def test_quarantine_s3_upload_when_file_exist(self, mock_s3_connector, file_upload): """Test to move file to quarantine bucket if exists.""" mock_s3_connector.object_exists.return_value = True mock_s3_connector.copy_object.return_value = { 'ResponseMetadata': {'HTTPStatusCode': 200} } mock_s3_connector.delete_object.return_value = { 'ResponseMetadata': {'HTTPStatusCode': 200} } quarantine_s3_upload(mock_s3_connector, file_upload) mock_s3_connector.object_exists.assert_called_once_with( 'test-bucket', 'test-key' ) mock_s3_connector.copy_object.assert_called_once_with( 'test-bucket', Config.S3_ABACUS_QUARANTINE_BUCKET, 'test-key' ) mock_s3_connector.delete_object.assert_called_once_with( 'test-bucket', 'test-key' ) def test_quarantine_s3_upload_when_file_not_exist( self, mock_s3_connector, file_upload ): """Test to move file to quarantine bucket if doesn't exist.""" mock_s3_connector.object_exists.return_value = False with pytest.raises(Exception) as exc_info: quarantine_s3_upload(mock_s3_connector, file_upload) assert ( f"File doesn't exist in source bucket {Config.S3_ABACUS_ADJUSTMENTS_BUCKET}" in str(exc_info.value) ) mock_s3_connector.object_exists.assert_called_once_with( 'test-bucket', 'test-key' ) mock_s3_connector.copy_object.assert_not_called() mock_s3_connector.delete_object.assert_not_called() def test_quarantine_s3_upload_copy_failure(self, mock_s3_connector, file_upload): """Test the error handler when the file copy operation fails.""" mock_s3_connector.object_exists.return_value = True mock_s3_connector.copy_object.side_effect = ClientError( {'Error': {'Code': '500', 'Message': 'Internal Error'}}, 'copy_object' ) with pytest.raises(ClientError) as exc_info: quarantine_s3_upload(mock_s3_connector, file_upload) assert ( 'An error occurred (500) when calling the copy_object operation: Internal Error' in str(exc_info.value) ) mock_s3_connector.copy_object.assert_called_once() mock_s3_connector.delete_object.assert_not_called() def test_quarantine_s3_upload_unexpected_copy_failure( self, mock_s3_connector, file_upload ): """Test the error handler when the file copy operation fails unexpectedly.""" mock_s3_connector.object_exists.return_value = True mock_s3_connector.copy_object.return_value = { 'ResponseMetadata': {'HTTPStatusCode': 503} } with pytest.raises(Exception) as exc_info: quarantine_s3_upload(mock_s3_connector, file_upload) assert f'Failed to copy object to {Config.S3_ABACUS_QUARANTINE_BUCKET}' in str( exc_info.value ) mock_s3_connector.copy_object.assert_called_once() mock_s3_connector.delete_object.assert_not_called() def test_quarantine_s3_upload_delete_failure(self, mock_s3_connector, file_upload): """Test the error handler when the file delete operation fails.""" mock_s3_connector.object_exists.return_value = True mock_s3_connector.copy_object.return_value = { 'ResponseMetadata': {'HTTPStatusCode': 200} } mock_s3_connector.delete_object.side_effect = ClientError( {'Error': {'Code': '500', 'Message': 'Internal Error'}}, 'delete_object' ) with pytest.raises(ClientError) as exc_info: quarantine_s3_upload(mock_s3_connector, file_upload) assert ( 'An error occurred (500) when calling the delete_object operation: Internal Error' in str(exc_info.value) ) mock_s3_connector.copy_object.assert_called_once() mock_s3_connector.delete_object.assert_called_once() def test_quarantine_s3_upload_unexpected_delete_failure( self, mock_s3_connector, file_upload ): """Test the error handler when the file delete operation fails unexpectedly.""" mock_s3_connector.object_exists.return_value = True mock_s3_connector.copy_object.return_value = { 'ResponseMetadata': {'HTTPStatusCode': 200} } mock_s3_connector.delete_object.return_value = { 'ResponseMetadata': {'HTTPStatusCode': 503} } with pytest.raises(Exception) as exc_info: quarantine_s3_upload(mock_s3_connector, file_upload) assert 'Failed to delete S3 object.' in str(exc_info.value) mock_s3_connector.copy_object.assert_called_once() mock_s3_connector.delete_object.assert_called_once()