如何对基于AWS STS与Athena的Python代码进行单元测试？

目标代码 #

以下是需要测试的Python代码（已补充遗漏的boto3导入）：

import boto3
import awswrangler as wr

def assume_role_and_query_athena(role_arn, query):
    # Assume the IAM role
    sts_client = boto3.client('sts')
    assumed_role = sts_client.assume_role(
        RoleArn=role_arn,
        RoleSessionName='AssumedRoleSession'
    )
    
    # Extract the temporary credentials from the AssumeRole response
    credentials = assumed_role['Credentials']
    access_key = credentials['AccessKeyId']
    secret_key = credentials['SecretAccessKey']
    session_token = credentials['SessionToken']
    
    # Use the temporary credentials to execute the query in Athena
    wr.athena.connect(
        access_key=access_key,
        secret_key=secret_key,
        session_token=session_token,
        region='us-east-1'  # Replace with your desired AWS region
    )
    df = wr.athena.read_sql_query(query)
    
    return df

可测试的核心点与实现方案 #

通过依赖模拟工具隔离真实AWS服务调用，就能覆盖代码的核心逻辑测试，具体可测试点包括：

• 验证STS角色假设的参数是否合规
• 验证临时凭证的提取逻辑是否正确
• 验证Athena连接与查询的参数传递是否准确
• 模拟异常场景（STS调用失败、Athena查询失败）验证代码表现

具体测试代码（基于pytest与unittest.mock） #

import pytest
from unittest.mock import patch, MagicMock
import pandas as pd
from your_module import assume_role_and_query_athena  # 替换为你的代码模块名

def test_assume_role_parameters():
    # 模拟STS客户端与返回结果
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.return_value = {
        'Credentials': {
            'AccessKeyId': 'test_access_key',
            'SecretAccessKey': 'test_secret_key',
            'SessionToken': 'test_session_token'
        }
    }
    
    with patch('boto3.client', return_value=mock_sts_client):
        with patch('awswrangler.athena.connect'):
            with patch('awswrangler.athena.read_sql_query', return_value=pd.DataFrame()):
                assume_role_and_query_athena('arn:aws:iam::123456789012:role/test-role', 'SELECT * FROM test_table')
                
                # 验证assume_role的参数是否正确
                mock_sts_client.assume_role.assert_called_once_with(
                    RoleArn='arn:aws:iam::123456789012:role/test-role',
                    RoleSessionName='AssumedRoleSession'
                )

def test_credentials_extraction_and_athena_connection():
    mock_sts_client = MagicMock()
    test_creds = {
        'AccessKeyId': 'test_access_key',
        'SecretAccessKey': 'test_secret_key',
        'SessionToken': 'test_session_token'
    }
    mock_sts_client.assume_role.return_value = {'Credentials': test_creds}
    
    mock_athena_connect = MagicMock()
    mock_read_query = MagicMock(return_value=pd.DataFrame())
    
    with patch('boto3.client', return_value=mock_sts_client):
        with patch('awswrangler.athena.connect', mock_athena_connect):
            with patch('awswrangler.athena.read_sql_query', mock_read_query):
                assume_role_and_query_athena('arn:aws:iam::123456789012:role/test-role', 'SELECT * FROM test_table')
                
                # 验证凭证是否正确传递给athena.connect
                mock_athena_connect.assert_called_once_with(
                    access_key='test_access_key',
                    secret_key='test_secret_key',
                    session_token='test_session_token',
                    region='us-east-1'
                )
                # 验证查询语句是否正确传递
                mock_read_query.assert_called_once_with('SELECT * FROM test_table')

def test_sts_assume_role_failure():
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.side_effect = Exception('STS role assume failed')
    
    with patch('boto3.client', return_value=mock_sts_client):
        with pytest.raises(Exception) as excinfo:
            assume_role_and_query_athena('arn:aws:iam::123456789012:role/test-role', 'SELECT * FROM test_table')
        
        assert 'STS role assume failed' in str(excinfo.value)

def test_athena_query_failure():
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.return_value = {
        'Credentials': {
            'AccessKeyId': 'test_access_key',
            'SecretAccessKey': 'test_secret_key',
            'SessionToken': 'test_session_token'
        }
    }
    
    with patch('boto3.client', return_value=mock_sts_client):
        with patch('awswrangler.athena.connect'):
            with patch('awswrangler.athena.read_sql_query', side_effect=Exception('Athena query failed')):
                with pytest.raises(Exception) as excinfo:
                    assume_role_and_query_athena('arn:aws:iam::123456789012:role/test-role', 'SELECT * FROM test_table')
                
                assert 'Athena query failed' in str(excinfo.value)

补充说明 #

• 运行测试前需安装依赖：pip install pytest pandas awswrangler boto3
• 所有测试均通过mock隔离外部依赖，不会真实调用AWS服务
• 若需贴近真实场景的集成测试，可使用moto库模拟AWS服务行为，但单元测试优先用unittest.mock保证速度与隔离性

内容的提问来源于stack exchange，提问作者koch