"""Tests for abacus_file_upload.logic.file_upload module."""

from datetime import datetime, timezone
from unittest.mock import MagicMock, Mock, patch

import pytest
from botocore.exceptions import ClientError
from marshmallow import ValidationError
from owsresponse import response

from abacus_file_upload.constants import (
    ABACUS_OUTBOX_EVENT_TYPES,
    ABACUS_OUTBOX_TARGET_TYPES,
    EVENT_PROCESSING_STATUSES,
    INVALID_STATUS_CHANGE,
    UPLOAD_STATUSES,
)
from abacus_file_upload.logic import file_upload


class TestInitiateUpload:
    """Tests for initiate_upload function."""

    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUploadConfig.find_by_upload_type')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.build')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_successful_single_part_upload(
        self,
        mock_commit,
        mock_build,
        mock_find_config,
        mock_get_s3,
        mock_file_upload_config,
        sample_upload_params,
    ):
        """Test successful single-part upload initiation."""
        # Setup mocks
        mock_find_config.return_value = mock_file_upload_config
        mock_s3 = Mock()
        mock_s3.generate_put_presigned_url.return_value = 'https://s3.amazonaws.com/...'
        mock_get_s3.return_value = mock_s3

        # Make file size small enough for single-part
        params = sample_upload_params.copy()
        params['file_size_bytes'] = 1024000  # 1MB

        response = file_upload.initiate_upload(**params)

        assert response.status == 201
        assert 'file_key' in response.message
        assert response.message['is_multipart'] is False
        assert 'upload_url' in response.message
        assert 'required_headers' in response.message
        assert 'Content-MD5' in response.message['required_headers']
        mock_commit.assert_called_once()

        # Verify S3 metadata includes filekey
        call_args = mock_s3.generate_put_presigned_url.call_args
        assert call_args is not None
        metadata = call_args[1]['metadata']
        assert 'filekey' in metadata
        assert metadata['filename'] == params['filename']
        assert metadata['uploadtype'] == params['upload_type']
        assert metadata['md5sum'] == params['md5sum']

    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUploadConfig.find_by_upload_type')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.build')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_successful_multipart_upload(
        self,
        mock_commit,
        mock_build,
        mock_find_config,
        mock_get_s3,
        mock_file_upload_config,
        sample_upload_params,
    ):
        """Test successful multipart upload initiation."""
        # Setup mocks
        # Increase max_file_size_bytes to allow 100MB file
        mock_file_upload_config.max_file_size_bytes = 1073741824  # 1GB
        mock_find_config.return_value = mock_file_upload_config

        mock_s3 = Mock()
        mock_s3.initiate_multipart_upload.return_value = 'test-upload-id'
        mock_s3.generate_multipart_presigned_urls.return_value = [
            {
                'part_number': 1,
                'url': 'https://s3.amazonaws.com/part1',
                'expires_at': datetime.now(timezone.utc),
            }
        ]
        mock_s3.generate_complete_multipart_presigned_url.return_value = (
            'https://s3.amazonaws.com/complete'
        )
        mock_get_s3.return_value = mock_s3

        # Make file size large enough for multipart
        params = sample_upload_params.copy()
        params['file_size_bytes'] = 104857600  # 100MB

        response = file_upload.initiate_upload(**params)

        assert response.status == 201
        assert 'file_key' in response.message
        assert response.message['is_multipart'] is True
        assert 'parts' in response.message
        assert 'complete_url' in response.message
        assert 'chunk_size_bytes' in response.message
        mock_commit.assert_called_once()

        # Verify S3 metadata includes filekey
        call_args = mock_s3.initiate_multipart_upload.call_args
        assert call_args is not None
        metadata = call_args[1]['metadata']
        assert 'filekey' in metadata
        assert metadata['filename'] == params['filename']
        assert metadata['uploadtype'] == params['upload_type']
        assert metadata['md5sum'] == params['md5sum']

    @patch('abacus_file_upload.logic.file_upload.FileUploadConfig.find_by_upload_type')
    def test_invalid_upload_type(self, mock_find_config, sample_upload_params):
        """Test initiate with invalid upload type."""
        mock_find_config.return_value = None

        response = file_upload.initiate_upload(**sample_upload_params)

        assert response.status == 404
        assert 'not found' in response.message.lower()

    @patch('abacus_file_upload.logic.file_upload.FileUploadConfig.find_by_upload_type')
    def test_invalid_file_type(self, mock_find_config, mock_file_upload_config):
        """Test initiate with invalid file type."""
        mock_find_config.return_value = mock_file_upload_config

        response = file_upload.initiate_upload(
            upload_type='adjustments',
            filename='test.pdf',  # PDF not allowed
            file_size_bytes=1024000,
            md5sum='5d41402abc4b2a76b9719d911017c592',
        )

        assert response.status == 400
        assert 'File type not allowed' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUploadConfig.find_by_upload_type')
    def test_file_size_exceeds_max(
        self, mock_find_config, mock_file_upload_config, sample_upload_params
    ):
        """Test initiate with file size exceeding maximum."""
        mock_find_config.return_value = mock_file_upload_config

        params = sample_upload_params.copy()
        params['file_size_bytes'] = 99999999999  # Way over 10MB limit

        response = file_upload.initiate_upload(**params)

        assert response.status == 400
        assert 'exceeds maximum allowed size' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUploadConfig.find_by_upload_type')
    def test_invalid_request_data(self, mock_find_config):
        """Test initiate with invalid request data."""
        response = file_upload.initiate_upload(
            upload_type='adjustments',
            filename='test.csv',
            # Missing required fields
        )

        assert response.status == 400

    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUploadConfig.find_by_upload_type')
    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    def test_s3_error(
        self, mock_rollback, mock_find_config, mock_get_s3, mock_file_upload_config
    ):
        """Test initiate with S3 error."""
        mock_find_config.return_value = mock_file_upload_config
        mock_s3 = Mock()
        mock_s3.generate_put_presigned_url.side_effect = ClientError(
            {'Error': {'Code': 'AccessDenied'}}, 'operation'
        )
        mock_get_s3.return_value = mock_s3

        response = file_upload.initiate_upload(
            upload_type='adjustments',
            filename='test.csv',
            file_size_bytes=1024000,
            md5sum='5d41402abc4b2a76b9719d911017c592',
        )

        assert response.status == 500
        assert 'Failed to initiate upload with S3' in response.message
        mock_rollback.assert_called_once()


class TestGetFileUpload:
    """Tests for get_file_upload function."""

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_schema')
    def test_get_existing_upload(self, mock_schema, mock_find, mock_file_upload):
        """Test getting existing file upload."""
        mock_find.return_value = mock_file_upload
        mock_schema.dump.return_value = {
            'file_key': '550e8400-e29b-41d4-a716-446655440000',
            'upload_status': 'init',
        }

        response = file_upload.get_file_upload('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 200
        assert 'file_key' in response.message
        mock_schema.dump.assert_called_once_with(mock_file_upload)

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_get_nonexistent_upload(self, mock_find):
        """Test getting nonexistent file upload."""
        mock_find.return_value = None

        response = file_upload.get_file_upload('nonexistent-key')

        assert response.status == 404
        assert 'Upload not found' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_verbose_schema')
    def test_get_existing_upload_verbose(
        self, mock_verbose_schema, mock_find, mock_file_upload
    ):
        """Test getting existing file upload with verbose=True."""
        mock_find.return_value = mock_file_upload
        mock_verbose_schema.dump.return_value = {
            'file_key': '550e8400-e29b-41d4-a716-446655440000',
            'upload_status': 'init',
            'created_at': '2025-11-24T12:00:00Z',
            'created_by': 'user123',
        }

        response = file_upload.get_file_upload(
            '550e8400-e29b-41d4-a716-446655440000', verbose=True
        )

        assert response.status == 200
        assert 'file_key' in response.message
        assert 'created_at' in response.message
        mock_verbose_schema.dump.assert_called_once_with(mock_file_upload)


class TestCompleteUpload:
    """Tests for complete_upload function."""

    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_schema')
    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.AbacusOutbox')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_successful_completion(
        self,
        mock_commit,
        mock_find,
        mock_abacus_outbox_model,
        mock_get_s3,
        mock_schema,
        mock_refresh,
    ):
        """Test successful upload completion."""
        mock_upload = Mock()
        mock_upload.file_key = '550e8400-e29b-41d4-a716-446655440000'
        mock_upload.upload_status = UPLOAD_STATUSES.INIT
        mock_upload.file_size_bytes = 1024000
        mock_upload.multipart_upload_id = None
        mock_upload.md5sum = '5d41402abc4b2a76b9719d911017c592'
        mock_upload.s3_bucket = 'test-bucket'
        mock_upload.s3_key = 'test-key'
        mock_upload.file_upload_id = 1
        mock_upload.file_upload_config = MagicMock(upload_type='flowthrough')
        mock_find.return_value = mock_upload
        mock_abacus_outbox_model.create.return_value = True

        mock_s3 = Mock()
        mock_s3.object_exists.return_value = True
        mock_s3.get_object_metadata.return_value = {
            'size': 1024000,
            'etag': '5d41402abc4b2a76b9719d911017c592',
        }
        mock_get_s3.return_value = mock_s3

        mock_schema.dump.return_value = {
            'file_key': '550e8400-e29b-41d4-a716-446655440000',
            'upload_status': 'complete',
        }

        resp = file_upload.complete_upload('550e8400-e29b-41d4-a716-446655440000')

        assert resp.status == 200
        mock_upload.update_attributes.assert_called_once()
        mock_commit.assert_called_once()
        mock_refresh.assert_called_once_with(mock_upload)
        mock_schema.dump.assert_called_once_with(mock_upload)
        mock_abacus_outbox_model.create.assert_called_once_with(
            target_type=ABACUS_OUTBOX_TARGET_TYPES.FILE_UPLOAD,
            target_id=mock_upload.file_upload_id,
            event_type=ABACUS_OUTBOX_EVENT_TYPES.FILE_UPLOAD_COMPLETED,
            status=EVENT_PROCESSING_STATUSES.PENDING,
            details={'upload_type': mock_upload.file_upload_config.upload_type},
        )

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_complete_nonexistent_upload(self, mock_find):
        """Test completing nonexistent upload."""
        mock_find.return_value = None

        response = file_upload.complete_upload('nonexistent-key')

        assert response.status == 404
        assert 'Upload not found' in response.message

    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_schema')
    @patch('abacus_file_upload.logic.file_upload.AbacusOutbox')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_complete_already_completed(
        self, mock_find, mock_abacus_outbox_model, mock_schema
    ):
        """Test completing already completed upload (idempotent)."""
        mock_upload = Mock()
        mock_upload.upload_status = UPLOAD_STATUSES.COMPLETE
        mock_find.return_value = mock_upload

        mock_schema.dump.return_value = {
            'file_key': '550e8400-e29b-41d4-a716-446655440000',
            'upload_status': 'complete',
        }

        resp = file_upload.complete_upload('550e8400-e29b-41d4-a716-446655440000')

        assert resp.status == 200
        mock_schema.dump.assert_called_once_with(mock_upload)
        mock_abacus_outbox_model.create.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.AbacusOutbox')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_complete_cancelled_upload(
        self, mock_find, mock_abacus_outbox_model, mock_file_upload
    ):
        """Test completing cancelled upload."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.CANCELLED
        mock_find.return_value = mock_file_upload

        response = file_upload.complete_upload('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 400
        assert 'Cannot complete a cancelled upload' in response.message
        mock_abacus_outbox_model.create.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.AbacusOutbox')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_file_not_in_s3(
        self,
        mock_commit,
        mock_find,
        mock_abacus_outbox_model,
        mock_get_s3,
        mock_file_upload,
    ):
        """Test completion when file doesn't exist in S3."""
        mock_find.return_value = mock_file_upload
        mock_s3 = Mock()
        mock_s3.object_exists.return_value = False
        mock_get_s3.return_value = mock_s3

        response = file_upload.complete_upload('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 400
        assert 'File not found in S3' in response.message
        mock_abacus_outbox_model.create.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_file_size_mismatch(
        self, mock_commit, mock_find, mock_get_s3, mock_file_upload
    ):
        """Test completion with file size mismatch."""
        mock_find.return_value = mock_file_upload
        mock_s3 = Mock()
        mock_s3.object_exists.return_value = True
        mock_s3.get_object_metadata.return_value = {
            'size': 999999,  # Different from expected 1024000
            'etag': '5d41402abc4b2a76b9719d911017c592',
        }
        mock_get_s3.return_value = mock_s3

        response = file_upload.complete_upload('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 400
        assert 'File size mismatch' in response.message


class TestCancelUpload:
    """Tests for cancel_upload function."""

    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_schema')
    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_cancel_multipart_upload(
        self, mock_commit, mock_find, mock_get_s3, mock_schema, mock_refresh
    ):
        """Test cancelling multipart upload."""
        mock_upload = Mock()
        mock_upload.multipart_upload_id = 'test-upload-id'
        mock_upload.upload_status = UPLOAD_STATUSES.INIT
        mock_upload.s3_bucket = 'test-bucket'
        mock_upload.s3_key = 'test-key'
        mock_find.return_value = mock_upload

        mock_s3 = Mock()
        mock_get_s3.return_value = mock_s3

        mock_schema.dump.return_value = {
            'file_key': '550e8400-e29b-41d4-a716-446655440000',
            'upload_status': 'cancelled',
        }

        resp = file_upload.cancel_upload('550e8400-e29b-41d4-a716-446655440000')

        assert resp.status == 200
        mock_s3.abort_multipart_upload.assert_called_once()
        mock_upload.update_attributes.assert_called_once()
        mock_commit.assert_called_once()
        mock_refresh.assert_called_once_with(mock_upload)
        mock_schema.dump.assert_called_once_with(mock_upload)

    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_schema')
    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_cancel_single_part_upload(
        self, mock_commit, mock_find, mock_get_s3, mock_schema, mock_refresh
    ):
        """Test cancelling single-part upload."""
        mock_upload = Mock()
        mock_upload.multipart_upload_id = None
        mock_upload.upload_status = UPLOAD_STATUSES.INIT
        mock_upload.s3_bucket = 'test-bucket'
        mock_upload.s3_key = 'test-key'
        mock_find.return_value = mock_upload

        mock_s3 = Mock()
        mock_s3.object_exists.return_value = True
        mock_get_s3.return_value = mock_s3

        mock_schema.dump.return_value = {
            'file_key': '550e8400-e29b-41d4-a716-446655440000',
            'upload_status': 'cancelled',
        }

        resp = file_upload.cancel_upload('550e8400-e29b-41d4-a716-446655440000')

        assert resp.status == 200
        mock_s3.delete_object.assert_called_once()
        mock_upload.update_attributes.assert_called_once()
        mock_commit.assert_called_once()
        mock_refresh.assert_called_once_with(mock_upload)
        mock_schema.dump.assert_called_once_with(mock_upload)

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_cancel_nonexistent_upload(self, mock_find):
        """Test cancelling nonexistent upload."""
        mock_find.return_value = None

        response = file_upload.cancel_upload('nonexistent-key')

        assert response.status == 404
        assert 'Upload not found' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_cancel_completed_upload(self, mock_find, mock_file_upload):
        """Test cancelling completed upload."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.COMPLETE
        mock_find.return_value = mock_file_upload

        response = file_upload.cancel_upload('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 400
        assert 'Cannot cancel a completed upload' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_cancel_already_cancelled(self, mock_find, mock_file_upload):
        """Test cancelling already cancelled upload (idempotent)."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.CANCELLED
        mock_find.return_value = mock_file_upload

        response = file_upload.cancel_upload('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 200


class TestDownloadFile:
    """Tests for download_file function."""

    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_successful_download_url_generation(
        self, mock_find, mock_get_s3, mock_file_upload
    ):
        """Test successful download URL generation."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.COMPLETE
        mock_find.return_value = mock_file_upload
        mock_s3 = Mock()
        mock_s3.generate_get_presigned_url.return_value = (
            'https://s3.amazonaws.com/download'
        )
        mock_get_s3.return_value = mock_s3

        response = file_upload.download_file('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 200
        assert 'download_url' in response.message
        assert 'file_name' in response.message
        assert 'file_size_bytes' in response.message
        mock_s3.generate_get_presigned_url.assert_called_once()

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_download_nonexistent_upload(self, mock_find):
        """Test downloading nonexistent upload."""
        mock_find.return_value = None

        response = file_upload.download_file('nonexistent-key')

        assert response.status == 404
        assert 'Upload not found' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_download_incomplete_upload(self, mock_find, mock_file_upload):
        """Test downloading incomplete upload."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.INIT
        mock_find.return_value = mock_file_upload

        response = file_upload.download_file('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 400
        assert 'not available for download' in response.message

    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_download_s3_error(self, mock_find, mock_get_s3, mock_file_upload):
        """Test download with S3 error."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.COMPLETE
        mock_find.return_value = mock_file_upload
        mock_s3 = Mock()
        mock_s3.generate_get_presigned_url.side_effect = ClientError(
            {'Error': {'Code': 'AccessDenied'}}, 'operation'
        )
        mock_get_s3.return_value = mock_s3

        response = file_upload.download_file('550e8400-e29b-41d4-a716-446655440000')

        assert response.status == 500
        assert 'Failed to generate download URL' in response.message


class TestQuarantineInfectedUpload:
    """Tests for quarantine_infected_upload function."""

    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_schema')
    @patch('abacus_file_upload.logic.file_upload.quarantine_s3_upload')
    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_quarantine_infected_upload(
        self,
        mock_commit,
        mock_find,
        mock_get_s3,
        mock_quarantine_s3_upload,
        mock_schema,
        mock_refresh,
    ):
        """Test moving file to another bucket."""
        mock_upload = Mock()
        mock_upload.upload_status = UPLOAD_STATUSES.INIT
        mock_upload.s3_bucket = 'test-bucket'
        mock_upload.s3_key = 'test-key'
        mock_find.return_value = mock_upload

        mock_s3 = Mock()
        mock_get_s3.return_value = mock_s3

        mock_schema.dump.return_value = {
            'file_key': '550e8400-e29b-41d4-a716-446655440000',
            'upload_status': 'quarantined',
        }

        resp = file_upload.quarantine_infected_upload(
            '550e8400-e29b-41d4-a716-446655440000'
        )

        assert resp.status == 200
        mock_quarantine_s3_upload.assert_called_once_with(mock_s3, mock_upload)
        mock_upload.update_attributes.assert_called_once()
        mock_commit.assert_called_once()
        mock_refresh.assert_called_once_with(mock_upload)
        mock_schema.dump.assert_called_once_with(mock_upload)

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_quarantine_nonexistent_upload(self, mock_find):
        """Test moving nonexistent upload."""
        mock_find.return_value = None

        response = file_upload.quarantine_infected_upload('nonexistent-key')

        assert response.status == 404
        assert 'Upload not found' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_quarantine_completed_upload(self, mock_find, mock_file_upload):
        """Test quarantine completed upload."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.COMPLETE
        mock_find.return_value = mock_file_upload

        response = file_upload.quarantine_infected_upload(
            '550e8400-e29b-41d4-a716-446655440000'
        )

        assert response.status == 400
        assert 'Cannot quarantine a completed upload' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_quarantine_cancelled_upload(self, mock_find, mock_file_upload):
        """Test quarantine cancelled upload."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.CANCELLED
        mock_find.return_value = mock_file_upload

        response = file_upload.quarantine_infected_upload(
            '550e8400-e29b-41d4-a716-446655440000'
        )

        assert response.status == 400
        assert 'Cannot quarantine a cancelled upload' in response.message

    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_already_quarantined(self, mock_find, mock_file_upload):
        """Test already quarantined upload."""
        mock_file_upload.upload_status = UPLOAD_STATUSES.QUARANTINED
        mock_find.return_value = mock_file_upload

        response = file_upload.quarantine_infected_upload(
            '550e8400-e29b-41d4-a716-446655440000'
        )

        assert response.status == 200

    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.file_upload_detail_schema')
    @patch('abacus_file_upload.logic.file_upload.quarantine_s3_upload')
    @patch('abacus_file_upload.logic.file_upload.get_s3_connector')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    def test_quarantine_infected_upload_error(
        self,
        mock_commit,
        mock_find,
        mock_get_s3,
        mock_quarantine_s3_upload,
        mock_schema,
        mock_refresh,
        mock_rollback,
    ):
        """Test quarantine_infected_upload with S3 error."""
        mock_upload = Mock()
        mock_upload.upload_status = UPLOAD_STATUSES.INIT
        mock_upload.s3_bucket = 'test-bucket'
        mock_upload.s3_key = 'test-key'
        mock_find.return_value = mock_upload

        mock_s3 = Mock()
        mock_get_s3.return_value = mock_s3

        mock_quarantine_s3_upload.side_effect = ClientError(
            {'Error': {'Code': 'AccessDenied'}}, 'operation'
        )

        response = file_upload.quarantine_infected_upload(
            '550e8400-e29b-41d4-a716-446655440000'
        )

        assert response.status == 500
        assert 'AccessDenied' in response.message

        mock_quarantine_s3_upload.assert_called_once_with(mock_s3, mock_upload)
        mock_upload.update_attributes.assert_called_once()
        mock_commit.assert_called_once()
        mock_refresh.assert_not_called()
        mock_schema.dump.assert_not_called()
        mock_rollback.assert_called_once()


class TestValidateUploadStatusTransition:
    """Tests for _validate_upload_status_transition function."""

    def test_transiting_to_same_status(self):
        """Test _validate_upload_status_transition when current and new status is same."""
        result = file_upload._validate_upload_status_transition(
            UPLOAD_STATUSES.COMPLETE, UPLOAD_STATUSES.COMPLETE
        )

        assert result is None

    def test_transiting_from_complete_to_quarantined(self):
        """Test _validate_upload_status_transition when updating from complete to quarantined."""
        with pytest.raises(ValidationError) as exec_info:
            file_upload._validate_upload_status_transition(
                UPLOAD_STATUSES.COMPLETE, UPLOAD_STATUSES.QUARANTINED
            )
        assert str(exec_info.value) == INVALID_STATUS_CHANGE.format(
            object_type='FileUpload',
            current_status=UPLOAD_STATUSES.COMPLETE,
            new_status=UPLOAD_STATUSES.QUARANTINED,
        )

    def test_transiting_from_scanning_to_complete(self):
        """Test _validate_upload_status_transition when updating from scanning to complete."""
        result = file_upload._validate_upload_status_transition(
            UPLOAD_STATUSES.SCANNING, UPLOAD_STATUSES.COMPLETE
        )
        result is None

    def test_transiting_from_scanning_to_error(self):
        """Test _validate_upload_status_transition when updating from scanning to error."""
        result = file_upload._validate_upload_status_transition(
            UPLOAD_STATUSES.SCANNING, UPLOAD_STATUSES.ERROR
        )

        assert result is None


class TestUpdateFileUploadStatus:
    """Tests for update_file_upload_status function."""

    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    @patch('abacus_file_upload.logic.file_upload._validate_upload_status_transition')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_update_file_upload_status(
        self,
        mock_find,
        mock_status_transition,
        mock_commit,
        mock_refresh,
        mock_rollback,
        mock_file_upload,
    ):
        """Test updating status of file_upload."""
        mock_find.return_value = mock_file_upload
        mock_status_transition.return_value = True
        mock_commit.return_value = True

        file_key = '550e8400-e29b-41d4-a716-446655440000'

        put_request_params = {'upload_status': 'scanning'}
        response = file_upload.update_file_upload_status(file_key, put_request_params)

        assert response.status == 200
        assert response.message['upload_status'] == 'scanning'
        mock_commit.assert_called_once()
        mock_refresh.assert_called_once()
        mock_rollback.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    @patch('abacus_file_upload.logic.file_upload._validate_upload_status_transition')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_non_existent_file_key(
        self,
        mock_find,
        mock_status_transition,
        mock_commit,
        mock_refresh,
        mock_rollback,
    ):
        """Test throws an error for non existent file_key."""
        mock_find.return_value = None
        file_key = '550e8400-e29b-41d4-a716-446655440000'

        put_request_params = {'upload_status': 'scanning'}
        response = file_upload.update_file_upload_status(file_key, put_request_params)

        assert response.status == 404
        assert response.message == f"Upload not found: '{file_key}'"
        mock_status_transition.assert_not_called()
        mock_commit.assert_not_called()
        mock_refresh.assert_not_called()
        mock_rollback.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    @patch('abacus_file_upload.logic.file_upload._validate_upload_status_transition')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_upload_status_transition_failed(
        self,
        mock_find,
        mock_status_transition,
        mock_commit,
        mock_refresh,
        mock_rollback,
        mock_file_upload,
    ):
        """Test throws an error for invalid upload_status transitioning.."""
        mock_find.return_value = mock_file_upload
        mock_status_transition.side_effect = ValidationError(
            'invalid batch status transition'
        )
        file_key = '550e8400-e29b-41d4-a716-446655440000'

        put_request_params = {'upload_status': 'complete'}
        response = file_upload.update_file_upload_status(file_key, put_request_params)

        assert response.status == 400
        assert response.message == 'invalid batch status transition'
        mock_status_transition.assert_called_once()
        mock_commit.assert_not_called()
        mock_refresh.assert_not_called()
        mock_rollback.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    @patch('abacus_file_upload.logic.file_upload.cancel_upload')
    @patch('abacus_file_upload.logic.file_upload._validate_upload_status_transition')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_call_cancel_upload_function(
        self,
        mock_find,
        mock_status_transition,
        mock_cancel_upload,
        mock_commit,
        mock_refresh,
        mock_rollback,
        mock_file_upload,
    ):
        """Test call cancel_upload function to update status to cancelled."""
        mock_find.return_value = mock_file_upload
        mock_status_transition.return_value = True
        file_key = '550e8400-e29b-41d4-a716-446655440000'
        mock_cancel_upload.return_value = response.Response(
            message={'upload_status': 'cancelled'}, status=200
        )

        put_request_params = {'upload_status': 'cancelled'}
        result = file_upload.update_file_upload_status(file_key, put_request_params)

        assert result.status == 200
        assert result.message['upload_status'] == 'cancelled'
        mock_status_transition.assert_called_once()
        mock_cancel_upload.assert_called_once_with(file_key)
        mock_commit.assert_not_called()
        mock_refresh.assert_not_called()
        mock_rollback.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    @patch('abacus_file_upload.logic.file_upload.complete_upload')
    @patch('abacus_file_upload.logic.file_upload._validate_upload_status_transition')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_call_complete_upload_function(
        self,
        mock_find,
        mock_status_transition,
        mock_complete_upload,
        mock_commit,
        mock_refresh,
        mock_rollback,
        mock_file_upload,
    ):
        """Test call complete_upload function to update status to complete."""
        mock_find.return_value = mock_file_upload
        mock_status_transition.return_value = True
        file_key = '550e8400-e29b-41d4-a716-446655440000'
        mock_complete_upload.return_value = response.Response(
            message={'upload_status': 'complete'}, status=200
        )

        put_request_params = {'upload_status': 'complete'}
        result = file_upload.update_file_upload_status(file_key, put_request_params)

        assert result.status == 200
        assert result.message['upload_status'] == 'complete'
        mock_status_transition.assert_called_once()
        mock_complete_upload.assert_called_once_with(file_key)
        mock_commit.assert_not_called()
        mock_refresh.assert_not_called()
        mock_rollback.assert_not_called()

    @patch('abacus_file_upload.logic.file_upload.db.session.rollback')
    @patch('abacus_file_upload.logic.file_upload.db.session.refresh')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.commit_changes')
    @patch('abacus_file_upload.logic.file_upload.quarantine_infected_upload')
    @patch('abacus_file_upload.logic.file_upload._validate_upload_status_transition')
    @patch('abacus_file_upload.logic.file_upload.FileUpload.find_by_file_key')
    def test_call_quarantined_upload_function(
        self,
        mock_find,
        mock_status_transition,
        mock_quarantine_infected_upload,
        mock_commit,
        mock_refresh,
        mock_rollback,
        mock_file_upload,
    ):
        """Test call quarantine_infected_upload function to update status to quarantined."""
        mock_find.return_value = mock_file_upload
        mock_status_transition.return_value = True
        file_key = '550e8400-e29b-41d4-a716-446655440000'
        mock_quarantine_infected_upload.return_value = response.Response(
            message={'upload_status': 'quarantined'}, status=200
        )

        put_request_params = {'upload_status': 'quarantined'}
        result = file_upload.update_file_upload_status(file_key, put_request_params)

        assert result.status == 200
        assert result.message['upload_status'] == 'quarantined'
        mock_status_transition.assert_called_once()
        mock_quarantine_infected_upload.assert_called_once_with(file_key)
        mock_commit.assert_not_called()
        mock_refresh.assert_not_called()
        mock_rollback.assert_not_called()