title: How to mock aws athena functionality.?
body: <p>I have an Athena class that has athena functionality. :</p>
  <pre><code>class _Athena:
      &quot;&quot;&quot;
      Creates connection and runs queries against
      AWS Athena.
      &quot;&quot;&quot;
  
      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):
              &quot;&quot;&quot;
              Gets Athena query details.
      
              :param query_id: id of athena query.
              :type query_id: str
              &quot;&quot;&quot;
              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 &quot;{}&quot; failed or'
                                      ' was cancelled'.format(self._query_string))
                  if status == 'SUCCEEDED':
                      location = response_get_query_details['QueryExecution'][
                          'ResultConfiguration']['OutputLocation']
                      LOGGER.info(&quot;Athena output location: %s&quot;, location)
                      return response_get_query_details
                  # if status of query is running or queued, wait
                  else:
                      time.sleep(5)
  </code></pre>
  <p>I want to mock <strong>_get_query_details</strong> as part of my unit test.  Here is the. function that I wrote to test it:</p>
  <pre><code>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)
  </code></pre>
  <p>However it fails with the following error :</p>
  <pre><code>AssertionError: {'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}} != &lt;MagicMock name='_Athena()._get_query_details()' id='4786490336'&gt;
  </code></pre>
  <p>In principle this test should pass based in the state is SUCCEEDED. Any ideas what I might be. doing wrong here.</p>