使用AWS Moto与Python Mock编写单元测试遇集群识别问题

因为你的collect_emr方法硬编码了通过assume_role创建boto3客户端的逻辑，又没法重构代码，核心思路是让方法内部创建的客户端自动接入Moto的mock环境，以下两种可行方案：

方案1：Mock boto3.client，强制返回Moto的EMR客户端 #

直接拦截boto3.client的调用，当方法内部创建EMR客户端时，返回Moto已经初始化好的mock客户端——不管传入的凭证是什么，都能访问到你在测试里创建的集群。

示例代码：

from unittest.mock import patch
import boto3
from moto import mock_aws

def test_collect_emr():
    with mock_aws():
        # 1. 在Moto环境创建测试用EMR集群
        emr_client = boto3.client('emr', region_name='us-east-1')
        cluster_id = emr_client.run_job_flow(
            Name='test-cluster',
            ReleaseLabel='emr-6.9.0',
            Instances={'InstanceGroups': [{'InstanceRole': 'MASTER', 'InstanceCount': 1, 'InstanceType': 'm5.xlarge'}]},
            Tags=[{'Key': 'env', 'Value': 'test'}]
        )['JobFlowId']
        cluster_arn = f"arn:aws:emr:us-east-1:123456789012:cluster/{cluster_id}"

        # 2. Mock boto3.client，让它返回Moto的EMR客户端
        def mock_boto_client(service, **kwargs):
            if service == 'emr':
                return emr_client
            # 其他服务正常创建客户端（如果有需要）
            return boto3.client(service, **kwargs)

        with patch('boto3.client', side_effect=mock_boto_client):
            # 3. 调用被测方法
            result = collect_emr()

            # 4. 断言结果符合预期
            assert len(result) == 1
            assert result[0]['arn'] == cluster_arn
            assert result[0]['tags'] == {'env': 'test'}

方案2：Mock STS的assume_role方法，返回Moto的临时凭证 #

拦截sts.assume_role的调用，让它返回Moto mock环境的临时凭证，这样collect_emr用这些凭证创建的客户端会自动接入Moto的mock集群。

示例代码：

from unittest.mock import patch
import boto3
from moto import mock_aws

def test_collect_emr():
    with mock_aws():
        # 1. 在Moto环境创建测试用EMR集群
        emr_client = boto3.client('emr', region_name='us-east-1')
        cluster_id = emr_client.run_job_flow(
            Name='test-cluster',
            ReleaseLabel='emr-6.9.0',
            Instances={'InstanceGroups': [{'InstanceRole': 'MASTER', 'InstanceCount': 1, 'InstanceType': 'm5.xlarge'}]},
            Tags=[{'Key': 'env', 'Value': 'test'}]
        )['JobFlowId']
        cluster_arn = f"arn:aws:emr:us-east-1:123456789012:cluster/{cluster_id}"

        # 2. 获取Moto的临时凭证（模拟assume_role返回的结果）
        sts_client = boto3.client('sts', region_name='us-east-1')
        mock_creds = sts_client.get_session_token()['Credentials']

        # 3. Mock sts.assume_role，返回Moto的凭证
        def mock_assume_role(**kwargs):
            return {
                'Credentials': {
                    'AccessKeyId': mock_creds['AccessKeyId'],
                    'SecretAccessKey': mock_creds['SecretAccessKey'],
                    'SessionToken': mock_creds['SessionToken'],
                    'Expiration': mock_creds['Expiration']
                }
            }

        # 注意：这里要patch collect_emr方法内部导入的STS模块路径，比如如果collect_emr里用boto3.client('sts')，就patch对应路径
        with patch('your_module.collect_emr.boto3.client') as mock_sts_client:
            mock_sts_client.return_value.assume_role.side_effect = mock_assume_role
            # 4. 调用被测方法
            result = collect_emr()

            # 5. 断言结果符合预期
            assert len(result) == 1
            assert result[0]['arn'] == cluster_arn
            assert result[0]['tags'] == {'env': 'test'}

关键注意事项 #

• 方案2中patch的路径必须是被测方法内部实际导入并使用的STS模块路径，比如如果collect_emr里是直接用boto3.client('sts')，就要patch该方法内部的boto3.client调用，确保拦截到assume_role请求。
• 两种方案都要确保Moto的mock_aws上下文在最外层，保证所有客户端都在mock环境内创建。

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