如何模拟 aws athena 功能。？

Question

我有一台 Athena class，它具有雅典娜功能。 :

class _Athena:
    """
    Creates connection and runs queries against
    AWS Athena.
    """

    def __init__(self, workgroup, database, params, query_string):
        self._workgroup = workgroup
        self._database = database
        self.params = params
        self._output_location = 's3://{bucket}/{path}'.format(
            bucket=_get_bucket(self.params),
            path=_get_path(self.params))
        self._client = _get_client(self.params)
        self._query_string = query_string

      def _get_query_details(self, query_id):
            """
            Gets Athena query details.
    
            :param query_id: id of athena query.
            :type query_id: str
            """
            while True:
                response_get_query_details = self._client.get_query_execution(
                    QueryExecutionId=query_id
                )
                status = response_get_query_details['QueryExecution'][
                    'Status']['State']
                LOGGER.info('Athena query status %s', status)
                if status in ('FAILED', 'CANCELLED'):
                    LOGGER.error(response_get_query_details)
                    raise Exception('Athena query with the string "{}" failed or'
                                    ' was cancelled'.format(self._query_string))
                if status == 'SUCCEEDED':
                    location = response_get_query_details['QueryExecution'][
                        'ResultConfiguration']['OutputLocation']
                    LOGGER.info("Athena output location: %s", location)
                    return response_get_query_details
                # if status of query is running or queued, wait
                else:
                    time.sleep(5)

我想模拟 _get_query_details 作为单元测试的一部分。这里是。我为测试它而编写的函数：

client = boto3.client('athena', 'us-east-1')


    @mock.patch.object(client, 'get_query_execution')
    @mock.patch('shared_utilities.verto_athena._Athena')
    def test_get_query_details(mock_class, mock_client):
        success = {'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}}
        mock_client.return_value = success
        result = mock_class.return_value._get_query_details('id')
        tools.assert_equal(success, result)

但是失败并出现以下错误：

AssertionError: {'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}} != <MagicMock name='_Athena()._get_query_details()' id='4786490336'>

原则上，该测试应该在状态为 SUCCEEDED 的情况下通过。任何想法我可能是什么。这里做错了。

Answer 1

我是通过以下方式实现的：

   params = {
        "region": "us-east-1",
        "bucket": "temp-prod",
        "path": "path/to/obj"
    }
    
    _athena = module._Athena('workgroup',
                             'database',
                              params,
                             'select * from query'
                             )
    
    
    
        @mock.patch.object(_athena._client, 'get_query_execution')
        def test_get_query_details_on_success(mock_client):
            success = {
                'QueryExecution':  {
                    'Status': {
                        'State': 'SUCCEEDED'
                    },
                    'ResultConfiguration':
                    {
                        'OutputLocation': 's3://output'
                    }
                }
            }
            mock_client.return_value = success
            result = _athena._get_query_details('id')
            tools.assert_equal(success, result)

如何模拟 aws athena 功能。？

How to mock aws athena functionality.?

mocking

nose

python-2.7