# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for the Finch MCP server."""

import os
import pytest
from awslabs.finch_mcp_server.consts import STATUS_ERROR, STATUS_SUCCESS
from awslabs.finch_mcp_server.server import (
    ensure_vm_running,
    finch_build_container_image,
    finch_create_ecr_repo,
    finch_push_image,
    sensitive_data_filter,
    set_enable_aws_resource_write,
)
from unittest.mock import MagicMock, patch


class TestSensitiveDataFilter:
    """Tests for the sensitive_data_filter function."""

    def test_filter_aws_access_key(self):
        """Test filtering AWS access keys."""
        record = {
            'message': 'AWS Access Key: AKIAIOSFODNN7EXAMPLE is sensitive'  # pragma: allowlist secret
        }

        sensitive_data_filter(record)

        assert 'AWS_ACCESS_KEY_REDACTED' in record['message']
        assert 'AKIAIOSFODNN7EXAMPLE' not in record['message']  # pragma: allowlist secret

    def test_filter_aws_secret_key(self):
        """Test filtering AWS secret keys."""
        record = {
            'message': 'AWS Secret Key: wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY is sensitive'
        }

        sensitive_data_filter(record)

        assert 'AWS_SECRET_KEY_REDACTED' in record['message']
        assert (
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY'  # pragma: allowlist secret
            not in record['message']
        )

    def test_filter_api_key(self):
        """Test filtering API keys."""
        record = {
            'message': "api_key='secret123'"  # pragma: allowlist secret
        }

        sensitive_data_filter(record)

        assert 'api_key=REDACTED' in record['message']
        assert 'secret123' not in record['message']

    def test_filter_password(self):
        """Test filtering passwords."""
        record = {
            'message': "password='mypassword'"  # pragma: allowlist secret
        }

        sensitive_data_filter(record)

        assert 'password=REDACTED' in record['message']
        assert 'mypassword' not in record['message']

    def test_filter_url_with_credentials(self):
        """Test filtering URLs with credentials."""
        record = {
            'message': 'Connection URL: https://username:password@example.com'  # pragma: allowlist secret
        }

        sensitive_data_filter(record)

        assert (
            'https://REDACTED:REDACTED@example.com'  # pragma: allowlist secret
            in record['message']
        )
        assert 'username:password' not in record['message']


class TestEnsureVmRunning:
    """Tests for the ensure_vm_running function."""

    @patch('awslabs.finch_mcp_server.server.get_vm_status')
    @patch('awslabs.finch_mcp_server.server.is_vm_nonexistent')
    @patch('awslabs.finch_mcp_server.server.is_vm_stopped')
    @patch('awslabs.finch_mcp_server.server.is_vm_running')
    @patch('awslabs.finch_mcp_server.server.initialize_vm')
    @patch('awslabs.finch_mcp_server.server.start_stopped_vm')
    @patch('awslabs.finch_mcp_server.server.format_result')
    @patch('sys.platform', 'darwin')
    def test_ensure_vm_running_already_running(
        self,
        mock_format_result,
        mock_start_vm,
        mock_initialize_vm,
        mock_is_running,
        mock_is_stopped,
        mock_is_nonexistent,
        mock_get_status,
    ):
        """Test ensure_vm_running function on macOS when VM is already running."""
        # VM is running
        mock_get_status.return_value = MagicMock()
        mock_is_nonexistent.return_value = False
        mock_is_stopped.return_value = False
        mock_is_running.return_value = True
        mock_format_result.return_value = {'status': STATUS_SUCCESS, 'message': 'VM is running'}

        result = ensure_vm_running()

        assert result['status'] == STATUS_SUCCESS
        mock_format_result.assert_called_with(STATUS_SUCCESS, 'Finch VM is already running.')
        mock_initialize_vm.assert_not_called()
        mock_start_vm.assert_not_called()

    @patch('awslabs.finch_mcp_server.server.get_vm_status')
    @patch('awslabs.finch_mcp_server.server.is_vm_nonexistent')
    @patch('awslabs.finch_mcp_server.server.is_vm_stopped')
    @patch('awslabs.finch_mcp_server.server.is_vm_running')
    @patch('awslabs.finch_mcp_server.server.initialize_vm')
    @patch('awslabs.finch_mcp_server.server.start_stopped_vm')
    @patch('awslabs.finch_mcp_server.server.format_result')
    @patch('sys.platform', 'darwin')
    def test_ensure_vm_running_stopped(
        self,
        mock_format_result,
        mock_start_vm,
        mock_initialize_vm,
        mock_is_running,
        mock_is_stopped,
        mock_is_nonexistent,
        mock_get_status,
    ):
        """Test ensure_vm_running function on macOS when VM is stopped."""
        # VM is stopped
        mock_get_status.return_value = MagicMock()
        mock_is_nonexistent.return_value = False
        mock_is_stopped.return_value = True
        mock_is_running.return_value = False
        mock_start_vm.return_value = {'status': STATUS_SUCCESS, 'message': 'VM started'}
        mock_format_result.return_value = {'status': STATUS_SUCCESS, 'message': 'VM started'}

        result = ensure_vm_running()

        assert result['status'] == STATUS_SUCCESS
        mock_start_vm.assert_called_once()
        mock_initialize_vm.assert_not_called()

    @patch('awslabs.finch_mcp_server.server.get_vm_status')
    @patch('awslabs.finch_mcp_server.server.is_vm_nonexistent')
    @patch('awslabs.finch_mcp_server.server.is_vm_stopped')
    @patch('awslabs.finch_mcp_server.server.is_vm_running')
    @patch('awslabs.finch_mcp_server.server.initialize_vm')
    @patch('awslabs.finch_mcp_server.server.start_stopped_vm')
    @patch('awslabs.finch_mcp_server.server.format_result')
    @patch('sys.platform', 'darwin')
    def test_ensure_vm_running_nonexistent(
        self,
        mock_format_result,
        mock_start_vm,
        mock_initialize_vm,
        mock_is_running,
        mock_is_stopped,
        mock_is_nonexistent,
        mock_get_status,
    ):
        """Test ensure_vm_running function on macOS when VM is nonexistent."""
        # VM is nonexistent
        mock_get_status.return_value = MagicMock()
        mock_is_nonexistent.return_value = True
        mock_is_stopped.return_value = False
        mock_is_running.return_value = False
        mock_initialize_vm.return_value = {'status': STATUS_SUCCESS, 'message': 'VM initialized'}
        mock_format_result.return_value = {'status': STATUS_SUCCESS, 'message': 'VM initialized'}

        result = ensure_vm_running()

        assert result['status'] == STATUS_SUCCESS
        mock_initialize_vm.assert_called_once()
        mock_start_vm.assert_not_called()

    @patch('awslabs.finch_mcp_server.server.get_vm_status')
    @patch('awslabs.finch_mcp_server.server.is_vm_nonexistent')
    @patch('awslabs.finch_mcp_server.server.is_vm_stopped')
    @patch('awslabs.finch_mcp_server.server.is_vm_running')
    @patch('awslabs.finch_mcp_server.server.initialize_vm')
    @patch('awslabs.finch_mcp_server.server.start_stopped_vm')
    @patch('awslabs.finch_mcp_server.server.format_result')
    @patch('sys.platform', 'darwin')  # Mock as macOS for testing
    def test_ensure_vm_running_failures(
        self,
        mock_format_result,
        mock_start_vm,
        mock_initialize_vm,
        mock_is_running,
        mock_is_stopped,
        mock_is_nonexistent,
        mock_get_status,
    ):
        """Test ensure_vm_running function on macOS when operations fail."""
        # Test VM start failure
        mock_get_status.return_value = MagicMock()
        mock_is_nonexistent.return_value = False
        mock_is_stopped.return_value = True
        mock_is_running.return_value = False
        mock_start_vm.return_value = {'status': STATUS_ERROR, 'message': 'Failed to start VM'}

        result = ensure_vm_running()

        assert result['status'] == STATUS_ERROR
        mock_start_vm.assert_called_once()
        mock_initialize_vm.assert_not_called()

        # Reset mocks for the next test
        mock_format_result.reset_mock()
        mock_initialize_vm.reset_mock()
        mock_start_vm.reset_mock()

        # Test VM initialization failure
        mock_is_nonexistent.return_value = True
        mock_is_stopped.return_value = False
        mock_is_running.return_value = False
        mock_initialize_vm.return_value = {
            'status': STATUS_ERROR,
            'message': 'Failed to initialize VM',
        }

        result = ensure_vm_running()

        assert result['status'] == STATUS_ERROR
        mock_initialize_vm.assert_called_once()
        mock_start_vm.assert_not_called()

    @patch('sys.platform', 'linux')
    @patch('awslabs.finch_mcp_server.server.format_result')
    def test_ensure_vm_running_on_linux(self, mock_format_result):
        """Test ensure_vm_running function on Linux."""
        mock_format_result.return_value = {
            'status': STATUS_SUCCESS,
            'message': 'Finch does not use a VM on Linux..',
        }

        result = ensure_vm_running()

        assert result['status'] == STATUS_SUCCESS
        assert result['message'] == 'Finch does not use a VM on Linux..'
        mock_format_result.assert_called_with(STATUS_SUCCESS, 'Finch does not use a VM on Linux..')


class TestFinchTools:
    """Tests for Finch operations in the server."""

    @pytest.mark.asyncio
    async def test_finch_build_container_image_success(self):
        """Test successful finch_build_container_image operation."""
        dockerfile_path = '/path/to/Dockerfile'
        context_path = '/path/to/context'
        tags = ['myimage:latest']
        platforms = ['linux/amd64']
        no_cache = False
        pull = True

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.contains_ecr_reference') as mock_contains_ecr,
            patch('awslabs.finch_mcp_server.server.configure_ecr') as mock_configure_ecr,
            patch('awslabs.finch_mcp_server.server.stop_vm') as mock_stop_vm,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
            patch('awslabs.finch_mcp_server.server.build_image') as mock_build_image,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_contains_ecr.return_value = False
            mock_ensure_vm.return_value = {'status': STATUS_SUCCESS}
            mock_build_image.return_value = {
                'status': STATUS_SUCCESS,
                'message': 'Successfully built image',
            }

            result = await finch_build_container_image(
                dockerfile_path=dockerfile_path,
                context_path=context_path,
                tags=tags,
                platforms=platforms,
                no_cache=no_cache,
                pull=pull,
            )

            assert result.status == STATUS_SUCCESS
            assert result.message == 'Successfully built image'

            mock_check_finch.assert_called_once()
            mock_contains_ecr.assert_called_once_with(dockerfile_path)
            mock_ensure_vm.assert_called_once()
            mock_build_image.assert_called_once()
            mock_configure_ecr.assert_not_called()
            mock_stop_vm.assert_not_called()

    @pytest.mark.asyncio
    async def test_finch_build_container_image_with_ecr(self):
        """Test finch_build_container_image with ECR reference."""
        dockerfile_path = '/path/to/Dockerfile'
        context_path = '/path/to/context'
        tags = ['123456789012.dkr.ecr.us-west-2.amazonaws.com/myrepo:latest']

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.contains_ecr_reference') as mock_contains_ecr,
            patch('awslabs.finch_mcp_server.server.configure_ecr') as mock_configure_ecr,
            patch('awslabs.finch_mcp_server.server.stop_vm') as mock_stop_vm,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
            patch('awslabs.finch_mcp_server.server.build_image') as mock_build_image,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_contains_ecr.return_value = True
            mock_configure_ecr.return_value = (
                {'status': STATUS_SUCCESS, 'message': 'Success'},
                True,
            )
            mock_ensure_vm.return_value = {'status': STATUS_SUCCESS}
            mock_build_image.return_value = {
                'status': STATUS_SUCCESS,
                'message': 'Successfully built image',
            }

            result = await finch_build_container_image(
                dockerfile_path=dockerfile_path,
                context_path=context_path,
                tags=tags,
            )

            assert result.status == STATUS_SUCCESS
            assert result.message == 'Successfully built image'

            mock_check_finch.assert_called_once()
            mock_contains_ecr.assert_called_once_with(dockerfile_path)
            mock_configure_ecr.assert_called_once()
            mock_stop_vm.assert_called_once_with(force=True)
            mock_ensure_vm.assert_called_once()
            mock_build_image.assert_called_once()

    @pytest.mark.asyncio
    async def test_finch_build_container_image_with_ecr_error(self):
        """Test finch_build_container_image with ECR reference when configure_ecr returns an error."""
        dockerfile_path = '/path/to/Dockerfile'
        context_path = '/path/to/context'
        tags = ['123456789012.dkr.ecr.us-west-2.amazonaws.com/myrepo:latest']

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.contains_ecr_reference') as mock_contains_ecr,
            patch('awslabs.finch_mcp_server.server.configure_ecr') as mock_configure_ecr,
            patch('awslabs.finch_mcp_server.server.stop_vm') as mock_stop_vm,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
            patch('awslabs.finch_mcp_server.server.build_image') as mock_build_image,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_contains_ecr.return_value = True
            mock_configure_ecr.return_value = (
                {'status': STATUS_ERROR, 'message': 'Failed to configure ECR'},
                False,
            )

            result = await finch_build_container_image(
                dockerfile_path=dockerfile_path,
                context_path=context_path,
                tags=tags,
            )

            assert result.status == STATUS_ERROR
            assert result.message == 'Failed to configure ECR'

            mock_check_finch.assert_called_once()
            mock_contains_ecr.assert_called_once_with(dockerfile_path)
            mock_configure_ecr.assert_called_once()
            mock_stop_vm.assert_not_called()
            mock_ensure_vm.assert_not_called()
            mock_build_image.assert_not_called()

    @pytest.mark.asyncio
    async def test_finch_build_container_image_finch_not_installed(self):
        """Test finch_build_container_image when Finch is not installed."""
        dockerfile_path = '/path/to/Dockerfile'
        context_path = '/path/to/context'

        with patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch:
            mock_check_finch.return_value = {
                'status': STATUS_ERROR,
                'message': 'Finch not installed',
            }

            result = await finch_build_container_image(
                dockerfile_path=dockerfile_path,
                context_path=context_path,
            )

            assert result.status == STATUS_ERROR
            assert result.message == 'Finch not installed'

            mock_check_finch.assert_called_once()

    @pytest.mark.asyncio
    async def test_finch_build_container_image_vm_error(self):
        """Test finch_build_container_image when VM fails to start."""
        dockerfile_path = '/path/to/Dockerfile'
        context_path = '/path/to/context'

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.contains_ecr_reference') as mock_contains_ecr,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_contains_ecr.return_value = False
            mock_ensure_vm.return_value = {'status': STATUS_ERROR, 'message': 'Failed to start VM'}

            result = await finch_build_container_image(
                dockerfile_path=dockerfile_path,
                context_path=context_path,
            )

            assert result.status == STATUS_ERROR
            assert result.message == 'Failed to start VM'

            mock_check_finch.assert_called_once()
            mock_contains_ecr.assert_called_once_with(dockerfile_path)
            mock_ensure_vm.assert_called_once()

    @pytest.mark.asyncio
    async def test_finch_build_container_image_build_error(self):
        """Test finch_build_container_image when build fails."""
        dockerfile_path = '/path/to/Dockerfile'
        context_path = '/path/to/context'

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.contains_ecr_reference') as mock_contains_ecr,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
            patch('awslabs.finch_mcp_server.server.build_image') as mock_build_image,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_contains_ecr.return_value = False
            mock_ensure_vm.return_value = {'status': STATUS_SUCCESS}
            mock_build_image.return_value = {'status': STATUS_ERROR, 'message': 'Build failed'}

            result = await finch_build_container_image(
                dockerfile_path=dockerfile_path,
                context_path=context_path,
            )

            assert result.status == STATUS_ERROR
            assert result.message == 'Build failed'

            mock_check_finch.assert_called_once()
            mock_contains_ecr.assert_called_once_with(dockerfile_path)
            mock_ensure_vm.assert_called_once()
            mock_build_image.assert_called_once()

    @pytest.mark.asyncio
    async def test_finch_build_container_image_exception(self):
        """Test finch_build_container_image when an exception occurs."""
        dockerfile_path = '/path/to/Dockerfile'
        context_path = '/path/to/context'

        with patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch:
            mock_check_finch.side_effect = Exception('Unexpected error')

            result = await finch_build_container_image(
                dockerfile_path=dockerfile_path,
                context_path=context_path,
            )

            assert result.status == STATUS_ERROR
            assert 'Error building Docker image: Unexpected error' in result.message

            mock_check_finch.assert_called_once()

    @pytest.mark.asyncio
    async def test_finch_push_image_success(self):
        """Test successful finch_push_image operation."""
        image = '123456789012.dkr.ecr.us-west-2.amazonaws.com/myrepo:latest'

        set_enable_aws_resource_write(True)

        try:
            with (
                patch(
                    'awslabs.finch_mcp_server.server.check_finch_installation'
                ) as mock_check_finch,
                patch('awslabs.finch_mcp_server.server.is_ecr_repository') as mock_is_ecr,
                patch('awslabs.finch_mcp_server.server.configure_ecr') as mock_configure_ecr,
                patch('awslabs.finch_mcp_server.server.stop_vm') as mock_stop_vm,
                patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
                patch('awslabs.finch_mcp_server.server.push_image') as mock_push_image,
            ):
                mock_check_finch.return_value = {'status': STATUS_SUCCESS}
                mock_is_ecr.return_value = True
                mock_configure_ecr.return_value = (
                    {'status': STATUS_SUCCESS, 'message': 'Success'},
                    True,
                )
                mock_ensure_vm.return_value = {'status': STATUS_SUCCESS}
                mock_push_image.return_value = {
                    'status': STATUS_SUCCESS,
                    'message': 'Successfully pushed image',
                }

                result = await finch_push_image(image=image)

                assert result.status == STATUS_SUCCESS
                assert result.message == 'Successfully pushed image'

                mock_check_finch.assert_called_once()
                mock_is_ecr.assert_called_once_with(image)
                mock_configure_ecr.assert_called_once()
                mock_stop_vm.assert_called_once_with(force=True)
                mock_ensure_vm.assert_called_once()
                mock_push_image.assert_called_once_with(image)
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_push_image_with_ecr_error(self):
        """Test finch_push_image with ECR reference when configure_ecr returns an error."""
        image = '123456789012.dkr.ecr.us-west-2.amazonaws.com/myrepo:latest'

        set_enable_aws_resource_write(True)

        try:
            with (
                patch(
                    'awslabs.finch_mcp_server.server.check_finch_installation'
                ) as mock_check_finch,
                patch('awslabs.finch_mcp_server.server.is_ecr_repository') as mock_is_ecr,
                patch('awslabs.finch_mcp_server.server.configure_ecr') as mock_configure_ecr,
                patch('awslabs.finch_mcp_server.server.stop_vm') as mock_stop_vm,
                patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
                patch('awslabs.finch_mcp_server.server.push_image') as mock_push_image,
            ):
                mock_check_finch.return_value = {'status': STATUS_SUCCESS}
                mock_is_ecr.return_value = True
                mock_configure_ecr.return_value = (
                    {'status': STATUS_ERROR, 'message': 'Failed to configure ECR'},
                    False,
                )

                result = await finch_push_image(image=image)

                assert result.status == STATUS_ERROR
                assert result.message == 'Failed to configure ECR'

                mock_check_finch.assert_called_once()
                mock_is_ecr.assert_called_once_with(image)
                mock_configure_ecr.assert_called_once()
                mock_stop_vm.assert_not_called()
                mock_ensure_vm.assert_not_called()
                mock_push_image.assert_not_called()
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_push_image_non_ecr(self):
        """Test finch_push_image with non-ECR repository."""
        image = 'docker.io/library/nginx:latest'

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.is_ecr_repository') as mock_is_ecr,
            patch('awslabs.finch_mcp_server.server.configure_ecr') as mock_configure_ecr,
            patch('awslabs.finch_mcp_server.server.stop_vm') as mock_stop_vm,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
            patch('awslabs.finch_mcp_server.server.push_image') as mock_push_image,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_is_ecr.return_value = False
            mock_ensure_vm.return_value = {'status': STATUS_SUCCESS}
            mock_push_image.return_value = {
                'status': STATUS_SUCCESS,
                'message': 'Successfully pushed image',
            }

            result = await finch_push_image(image=image)

            assert result.status == STATUS_SUCCESS
            assert result.message == 'Successfully pushed image'

            mock_check_finch.assert_called_once()
            mock_is_ecr.assert_called_once_with(image)
            mock_configure_ecr.assert_not_called()
            mock_stop_vm.assert_not_called()
            mock_ensure_vm.assert_called_once()
            mock_push_image.assert_called_once_with(image)

    @pytest.mark.asyncio
    async def test_finch_push_image_finch_not_installed(self):
        """Test finch_push_image when Finch is not installed."""
        image = 'docker.io/library/nginx:latest'

        with patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch:
            mock_check_finch.return_value = {
                'status': STATUS_ERROR,
                'message': 'Finch not installed',
            }

            result = await finch_push_image(image=image)

            assert result.status == STATUS_ERROR
            assert result.message == 'Finch not installed'

            mock_check_finch.assert_called_once()

    @pytest.mark.asyncio
    async def test_finch_push_image_vm_error(self):
        """Test finch_push_image when VM fails to start."""
        image = 'docker.io/library/nginx:latest'

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.is_ecr_repository') as mock_is_ecr,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_is_ecr.return_value = False
            mock_ensure_vm.return_value = {'status': STATUS_ERROR, 'message': 'Failed to start VM'}

            result = await finch_push_image(image=image)

            assert result.status == STATUS_ERROR
            assert result.message == 'Failed to start VM'

            mock_check_finch.assert_called_once()
            mock_is_ecr.assert_called_once_with(image)
            mock_ensure_vm.assert_called_once()

    @pytest.mark.asyncio
    async def test_finch_push_image_push_error(self):
        """Test finch_push_image when push fails."""
        image = 'docker.io/library/nginx:latest'
        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.is_ecr_repository') as mock_is_ecr,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
            patch('awslabs.finch_mcp_server.server.push_image') as mock_push_image,
        ):
            mock_check_finch.return_value = {'status': STATUS_SUCCESS}
            mock_is_ecr.return_value = False
            mock_ensure_vm.return_value = {'status': STATUS_SUCCESS}
            mock_push_image.return_value = {'status': STATUS_ERROR, 'message': 'Push failed'}

            result = await finch_push_image(image=image)

            assert result.status == STATUS_ERROR
            assert result.message == 'Push failed'

            mock_check_finch.assert_called_once()
            mock_is_ecr.assert_called_once_with(image)
            mock_ensure_vm.assert_called_once()
            mock_push_image.assert_called_once_with(image)

    @pytest.mark.asyncio
    async def test_finch_push_image_exception(self):
        """Test finch_push_image when an exception occurs."""
        image = 'docker.io/library/nginx:latest'

        set_enable_aws_resource_write(True)

        try:
            with patch(
                'awslabs.finch_mcp_server.server.check_finch_installation'
            ) as mock_check_finch:
                mock_check_finch.side_effect = Exception('Unexpected error')

                result = await finch_push_image(image=image)

                assert result.status == STATUS_ERROR
                assert 'Error pushing image: Unexpected error' in result.message

                mock_check_finch.assert_called_once()
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_create_ecr_repo_success(self):
        """Test successful finch_create_ecr_repo operation."""
        repository_name = 'test-repo'
        region = 'us-west-2'

        set_enable_aws_resource_write(True)

        try:
            with patch(
                'awslabs.finch_mcp_server.server.create_ecr_repository'
            ) as mock_create_ecr_repository:
                mock_create_ecr_repository.return_value = {
                    'status': STATUS_SUCCESS,
                    'message': "Successfully created ECR repository 'test-repo'.",
                }

                result = await finch_create_ecr_repo(
                    repository_name=repository_name, region=region
                )

                assert result.status == STATUS_SUCCESS
                assert "Successfully created ECR repository 'test-repo'" in result.message

                mock_create_ecr_repository.assert_called_once_with(
                    repository_name=repository_name, region=region
                )
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_create_ecr_repo_already_exists(self):
        """Test finch_create_ecr_repo when repository already exists."""
        repository_name = 'test-repo'
        region = 'us-west-2'

        set_enable_aws_resource_write(True)

        try:
            with patch(
                'awslabs.finch_mcp_server.server.create_ecr_repository'
            ) as mock_create_ecr_repository:
                mock_create_ecr_repository.return_value = {
                    'status': STATUS_SUCCESS,
                    'message': "ECR repository 'test-repo' already exists.",
                    'repository_uri': '123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo',
                    'exists': True,
                }

                result = await finch_create_ecr_repo(
                    repository_name=repository_name, region=region
                )

                assert result.status == STATUS_SUCCESS
                assert 'already exists' in result.message

                mock_create_ecr_repository.assert_called_once_with(
                    repository_name=repository_name, region=region
                )
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_create_ecr_repo_error(self):
        """Test finch_create_ecr_repo when creation fails."""
        repository_name = 'test-repo'
        region = 'us-west-2'

        set_enable_aws_resource_write(True)

        try:
            with patch(
                'awslabs.finch_mcp_server.server.create_ecr_repository'
            ) as mock_create_ecr_repository:
                mock_create_ecr_repository.return_value = {
                    'status': STATUS_ERROR,
                    'message': "Failed to create ECR repository 'test-repo': Access denied",
                }

                result = await finch_create_ecr_repo(
                    repository_name=repository_name, region=region
                )

                assert result.status == STATUS_ERROR
                assert 'Failed to create ECR repository' in result.message

                mock_create_ecr_repository.assert_called_once_with(
                    repository_name=repository_name, region=region
                )
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_create_ecr_repo_exception(self):
        """Test finch_create_ecr_repo when an exception occurs."""
        repository_name = 'test-repo'
        region = 'us-west-2'

        set_enable_aws_resource_write(True)

        try:
            with patch(
                'awslabs.finch_mcp_server.server.create_ecr_repository'
            ) as mock_create_ecr_repository:
                mock_create_ecr_repository.side_effect = Exception('Unexpected error')

                result = await finch_create_ecr_repo(
                    repository_name=repository_name, region=region
                )

                assert result.status == STATUS_ERROR
                assert 'Error checking/creating ECR repository: Unexpected error' in result.message

                mock_create_ecr_repository.assert_called_once_with(
                    repository_name=repository_name, region=region
                )
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_create_ecr_repo_readonly_mode(self):
        """Test finch_create_ecr_repo when AWS resource write is disabled (which is the default)."""
        repository_name = 'test-repo'
        region = 'us-west-2'

        try:
            result = await finch_create_ecr_repo(repository_name=repository_name, region=region)

            assert result.status == STATUS_ERROR
            assert (
                result.message == 'Server running in read-only mode, unable to perform the action'
            )
        finally:
            set_enable_aws_resource_write(False)

    @pytest.mark.asyncio
    async def test_finch_push_image_readonly_mode(self):
        """Test finch_push_image when AWS resource write is disabled and pushing to ECR."""
        image = '123456789012.dkr.ecr.us-west-2.amazonaws.com/myrepo:latest'

        with (
            patch('awslabs.finch_mcp_server.server.check_finch_installation') as mock_check_finch,
            patch('awslabs.finch_mcp_server.server.is_ecr_repository') as mock_is_ecr,
            patch('awslabs.finch_mcp_server.server.configure_ecr') as mock_configure_ecr,
            patch('awslabs.finch_mcp_server.server.stop_vm') as mock_stop_vm,
            patch('awslabs.finch_mcp_server.server.ensure_vm_running') as mock_ensure_vm,
            patch('awslabs.finch_mcp_server.server.push_image') as mock_push_image,
        ):
            mock_check_finch.return_value = {'status': 'success'}
            mock_is_ecr.return_value = True

            try:
                result = await finch_push_image(image=image)

                assert result.status == STATUS_ERROR
                assert (
                    result.message
                    == 'Server running in read-only mode, unable to push to ECR repository'
                )
                mock_check_finch.assert_called_once()
                mock_is_ecr.assert_called_once_with(image)
                mock_configure_ecr.assert_not_called()
                mock_stop_vm.assert_not_called()
                mock_ensure_vm.assert_not_called()
                mock_push_image.assert_not_called()
            finally:
                set_enable_aws_resource_write(False)


class TestMainFunctionLogging:
    """Tests for main function logging behavior."""

    @patch.dict(os.environ, {'FINCH_MCP_LOG_FILE': '/custom/test.log'})
    @patch('awslabs.finch_mcp_server.server.logger')
    @patch('awslabs.finch_mcp_server.server.mcp')
    def test_main_logs_custom_file_path(self, mock_mcp, mock_logger):
        """Test main function logs custom file path correctly."""
        from awslabs.finch_mcp_server.server import main

        main()
        mock_logger.info.assert_any_call('Logging to stderr and file: /custom/test.log')

    @patch.dict(os.environ, {'FINCH_DISABLE_FILE_LOGGING': 'true'})
    @patch('awslabs.finch_mcp_server.server.logger')
    @patch('awslabs.finch_mcp_server.server.mcp')
    def test_main_logs_stderr_only(self, mock_mcp, mock_logger):
        """Test main function logs stderr-only mode correctly."""
        from awslabs.finch_mcp_server.server import main

        main()
        mock_logger.warning.assert_any_call('Logging to stderr only')

    @patch('awslabs.finch_mcp_server.server.logger')
    @patch('awslabs.finch_mcp_server.server.mcp')
    def test_main_logs_default_logging(self, mock_mcp, mock_logger):
        """Test main function logs default logging mode correctly."""
        from awslabs.finch_mcp_server.server import main

        main()
        mock_logger.info.assert_any_call('Logging to stderr and default logging file')
