from unittest.mock import MagicMock, mock_open, patch

import pytest
import yaml

from litellm.proxy.common_utils.load_config_utils import get_file_contents_from_s3


class TestGetFileContentsFromS3:
    """Test suite for S3 config loading functionality."""

    @patch('boto3.client')
    @patch('litellm.main.bedrock_converse_chat_completion')
    @patch('yaml.safe_load')
    def test_get_file_contents_from_s3_no_temp_file_creation(
        self, mock_yaml_load, mock_bedrock, mock_boto3_client
    ):
        """
        Test that get_file_contents_from_s3 doesn't create temporary files
        and uses yaml.safe_load directly on the S3 response content.

        Note: It's critical that yaml.safe_load is used

        Relevant issue/PR: https://github.com/BerriAI/litellm/pull/12078
        """
        # Mock credentials
        mock_credentials = MagicMock()
        mock_credentials.access_key = "test_access_key"
        mock_credentials.secret_key = "test_secret_key"
        mock_credentials.token = "test_token"
        mock_bedrock.get_credentials.return_value = mock_credentials

        # Mock S3 client and response
        mock_s3_client = MagicMock()
        mock_boto3_client.return_value = mock_s3_client
        
        # Mock S3 response with YAML content
        yaml_content = """
        model_list:
          - model_name: gpt-3.5-turbo
            litellm_params:
              model: gpt-3.5-turbo
        """
        mock_response_body = MagicMock()
        mock_response_body.read.return_value = yaml_content.encode('utf-8')
        mock_s3_response = {
            'Body': mock_response_body
        }
        mock_s3_client.get_object.return_value = mock_s3_response

        # Mock yaml.safe_load to return parsed config
        expected_config = {
            'model_list': [{
                'model_name': 'gpt-3.5-turbo',
                'litellm_params': {
                    'model': 'gpt-3.5-turbo'
                }
            }]
        }
        mock_yaml_load.return_value = expected_config

        # Call the function
        bucket_name = "test-bucket"
        object_key = "config.yaml"
        result = get_file_contents_from_s3(bucket_name, object_key)

        # Assertions
        assert result == expected_config
        
        # Verify S3 client was created with correct credentials
        mock_boto3_client.assert_called_once_with(
            "s3",
            aws_access_key_id="test_access_key",
            aws_secret_access_key="test_secret_key",
            aws_session_token="test_token"
        )
        
        # Verify S3 get_object was called with correct parameters
        mock_s3_client.get_object.assert_called_once_with(
            Bucket=bucket_name,
            Key=object_key
        )
        
        # Verify the response body was read and decoded
        mock_response_body.read.assert_called_once()
        
        # Verify yaml.safe_load was called with the decoded content
        mock_yaml_load.assert_called_once_with(yaml_content)