# 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 profile option in server.py."""

import awslabs.aurora_dsql_mcp_server.server
from awslabs.aurora_dsql_mcp_server.server import main, _config
from unittest.mock import patch


class TestProfileOption:
    """Tests for the profile option."""

    @patch(
        'sys.argv',
        [
            'awslabs.aurora-dsql-mcp-server',
            '--cluster_endpoint',
            'test_ce',
            '--database_user',
            'test_user',
            '--region',
            'us-west-2',
            '--profile',
            'test-profile',
        ],
    )
    def test_main_with_profile_argument(self, mocker):
        mock_boto3_session = mocker.patch('boto3.Session')
        mock_session = mock_boto3_session.return_value
        mock_dsql_client = mock_session.client.return_value

        mock_mcp_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.mcp.run')

        main()

        # Check that the profile was set correctly
        assert awslabs.aurora_dsql_mcp_server.server.aws_profile == 'test-profile'

        # Check that boto3.Session was called with the correct profile
        mock_boto3_session.assert_called_once_with(profile_name='test-profile')

        # Check that the session's client method was called with the correct service, region, and config
        mock_session.client.assert_called_once_with('dsql', region_name='us-west-2', config=_config)

        # Check that the dsql client was set correctly
        assert awslabs.aurora_dsql_mcp_server.server.dsql_client == mock_dsql_client

        # Check that the server was started
        mock_mcp_run.assert_called_once()

    @patch(
        'sys.argv',
        [
            'awslabs.aurora-dsql-mcp-server',
            '--cluster_endpoint',
            'test_ce',
            '--database_user',
            'test_user',
            '--region',
            'us-west-2',
        ],
    )
    def test_main_without_profile_argument(self, mocker):
        mock_boto3_session = mocker.patch('boto3.Session')
        mock_session = mock_boto3_session.return_value
        mock_dsql_client = mock_session.client.return_value

        mock_mcp_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.mcp.run')

        main()

        # Check that the profile was not set
        assert awslabs.aurora_dsql_mcp_server.server.aws_profile is None

        # Check that boto3.Session was called without a profile
        mock_boto3_session.assert_called_once_with()

        # Check that the session's client method was called with the correct service, region, and config
        mock_session.client.assert_called_once_with('dsql', region_name='us-west-2', config=_config)

        # Check that the dsql client was set correctly
        assert awslabs.aurora_dsql_mcp_server.server.dsql_client == mock_dsql_client

        # Check that the server was started
        mock_mcp_run.assert_called_once()
