Question

我正在使用unittest来测试使用boto3调用AWS的函数。

该功能如下所示：

import boto3


def my_function():
    client = boto3.client('athena')
    res = client.start_query_exeuction(
        QueryString='SELECT * FROM logs',
        ResultConfiguration={'OutputLocation': 's3://mybucket'}
    )

    return res['QueryExecutionId']

我正在使用botocore stubber在我的单元测试中存根这个请求：

from botocore.stub import Stubber
import botocore.session

def test_my_function():    
    client = botocore.session.get_session().create_client('athena')
    client_res = {'QueryExecutionId': 'testid'}
    exp_params = {
        'QueryString': 'SELECT * FROM logs',
        'ResultConfiguration': {
            'OutputLocation': 's3://mybucket'
        }
    }
    with Stubber(client) as stubber:
        stubber.add_response('start_query_execution', client_res, exp_params)
        res = my_function()

    self.assertEqual(res, 'testid')

此测试失败

botocore.exceptions.ClientError：发生错误（UnrecognizedClientException）在调用StartQueryExecution时 operation：请求中包含的安全令牌无效。

为什么会失败？是因为我在my_function()中创建了一个与客户端中使用的客户端不同的新客户端吗？如果是这样，我该怎么测试呢？

非常感谢任何帮助。

Answer 1

目前，my_function()正在创建一个新客户端，并使用该客户端代替stubber。

一种选择是改变my_function以_client为参数。

def my_function(_client=None):
    if _client is not None:
        client = _client
    else:
        client = boto3.client('athena')
    res = client.start_query_exeuction(
        QueryString='SELECT * FROM logs',
        ResultConfiguration={'OutputLocation': 's3://mybucket'}
    )

    return res['QueryExecutionId']

然后将stubber传递给my_function。

with Stubber(client) as stubber:
    stubber.add_response('start_query_execution', client_res, exp_params)
    res = my_function(_client=stubber)

另一种选择是使用mock来修补boto.client以返回您的存根。

Answer 2

您还可以为客户端命名空间并执行以下操作：

mymodule.py

import boto3

class Amazon
    client = boto3.client('athena') # giving easy namespace access
    @classmethod
    def my_function(cls):
        res = cls.client.start_query_exeuction(
                QueryString='SELECT * FROM logs',
                ResultConfiguration={'OutputLocation': 's3://mybucket'}
            )

        return res['QueryExecutionId']

然后在您的测试中执行以下操作：

testmymodule.py

from botocore.stub import Stubber
from mymodule import Amazon

def test_my_function():
    client_res = {'QueryExecutionId': 'testid'}
    exp_params = {
        'QueryString': 'SELECT * FROM logs',
        'ResultConfiguration': {
            'OutputLocation': 's3://mybucket'
        }
    }
    with Stubber(Amazon.client) as stubber:
        stubber.add_response('start_query_execution', client_res, exp_params)
        res = Amazon.my_function()

    self.assertEqual(res, 'testid')

使用botocore stubber时出现UnrecognizedClientException

2 个答案: