如何在Python中Mock Athena并对批量处理相关代码做单元测试？

Absolutely! You can mock Athena in Python to write reliable unit tests for your batch processing functions without hitting real AWS resources. The two most common approaches are using botocore.stub.Stubber (purpose-built for AWS SDK interactions) or the standard unittest.mock library. Let’s walk through both with examples tailored to your Athena use case.

1. Using botocore.stub.Stubber (Recommended for AWS Services) #

The Stubber from botocore lets you predefine responses to Athena API calls and validate that your code uses the correct parameters. This is great for ensuring your function interacts with Athena exactly as intended.

Example Function to Test #

First, let’s assume you have a function that submits an Athena query, waits for completion, and returns formatted results:

import boto3
import time

def run_athena_query(query: str, database: str, output_location: str) -> list:
    athena_client = boto3.client('athena')
    
    # Submit query
    response = athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={'Database': database},
        ResultConfiguration={'OutputLocation': output_location}
    )
    query_execution_id = response['QueryExecutionId']
    
    # Wait for query completion
    while True:
        status_response = athena_client.get_query_execution(QueryExecutionId=query_execution_id)
        status = status_response['QueryExecution']['Status']['State']
        if status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
            break
        time.sleep(1)
    
    if status != 'SUCCEEDED':
        raise Exception(f"Query failed with status {status}")
    
    # Fetch and format results
    results_response = athena_client.get_query_results(QueryExecutionId=query_execution_id)
    header = [col['VarCharValue'] for col in results_response['ResultSet']['Rows'][0]['Data']]
    rows = []
    for row in results_response['ResultSet']['Rows'][1:]:
        rows.append(dict(zip(header, [col.get('VarCharValue') for col in row['Data']])))
    
    return rows

Unit Test with Stubber #

Here’s how to test this function using Stubber, covering both success and failure scenarios:

import unittest
from botocore.stub import Stubber
import boto3
from your_module import run_athena_query
from unittest.mock import patch

class TestAthenaBatchFunctions(unittest.TestCase):
    def test_run_athena_query_success(self):
        # Initialize Athena client and stubber
        athena_client = boto3.client('athena')
        with Stubber(athena_client) as stubber:
            # Stub query submission response
            stubber.add_response(
                'start_query_execution',
                {'QueryExecutionId': 'test-query-123'},
                expected_params={
                    'QueryString': 'SELECT * FROM user_data',
                    'QueryExecutionContext': {'Database': 'analytics_db'},
                    'ResultConfiguration': {'OutputLocation': 's3://my-bucket/athena-results/'}
                }
            )
            
            # Stub status checks (first RUNNING, then SUCCEEDED)
            stubber.add_response(
                'get_query_execution',
                {'QueryExecution': {'Status': {'State': 'RUNNING'}}},
                expected_params={'QueryExecutionId': 'test-query-123'}
            )
            stubber.add_response(
                'get_query_execution',
                {'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}},
                expected_params={'QueryExecutionId': 'test-query-123'}
            )
            
            # Stub results response
            stubber.add_response(
                'get_query_results',
                {
                    'ResultSet': {
                        'Rows': [
                            {'Data': [{'VarCharValue': 'user_id'}, {'VarCharValue': 'email'}]},
                            {'Data': [{'VarCharValue': '101'}, {'VarCharValue': 'alice@example.com'}]},
                            {'Data': [{'VarCharValue': '102'}, {'VarCharValue': 'bob@example.com'}]}
                        ]
                    }
                },
                expected_params={'QueryExecutionId': 'test-query-123'}
            )
            
            # Patch boto3.client to return our stubbed client
            with patch('boto3.client', return_value=athena_client):
                result = run_athena_query(
                    query='SELECT * FROM user_data',
                    database='analytics_db',
                    output_location='s3://my-bucket/athena-results/'
                )
                
                # Validate results
                self.assertEqual(len(result), 2)
                self.assertEqual(result[0]['email'], 'alice@example.com')
                # Ensure all stubbed calls were used
                stubber.assert_no_pending_responses()

    def test_run_athena_query_failure(self):
        athena_client = boto3.client('athena')
        with Stubber(athena_client) as stubber:
            # Stub query submission
            stubber.add_response(
                'start_query_execution',
                {'QueryExecutionId': 'test-query-456'},
                expected_params={
                    'QueryString': 'SELECT * FROM invalid_table',
                    'QueryExecutionContext': {'Database': 'analytics_db'},
                    'ResultConfiguration': {'OutputLocation': 's3://my-bucket/athena-results/'}
                }
            )
            
            # Stub failed status
            stubber.add_response(
                'get_query_execution',
                {
                    'QueryExecution': {
                        'Status': {'State': 'FAILED', 'StateChangeReason': 'Table not found'}
                    }
                },
                expected_params={'QueryExecutionId': 'test-query-456'}
            )
            
            with patch('boto3.client', return_value=athena_client):
                with self.assertRaises(Exception) as context:
                    run_athena_query(
                        query='SELECT * FROM invalid_table',
                        database='analytics_db',
                        output_location='s3://my-bucket/athena-results/'
                    )
                
                self.assertIn('Query failed with status FAILED', str(context.exception))
                stubber.assert_no_pending_responses()

if __name__ == '__main__':
    unittest.main()

2. Using unittest.mock.patch #

If you prefer a more flexible approach (without strict parameter validation), you can use Python’s built-in unittest.mock to mock the Athena client directly.

Example Test with MagicMock #

import unittest
from unittest.mock import patch, MagicMock
from your_module import run_athena_query

class TestAthenaFunctionsWithMock(unittest.TestCase):
    def test_run_athena_query_success(self):
        # Create a mock Athena client
        mock_athena = MagicMock()
        
        # Define mock responses
        mock_athena.start_query_execution.return_value = {'QueryExecutionId': 'test-query-789'}
        mock_athena.get_query_execution.side_effect = [
            {'QueryExecution': {'Status': {'State': 'RUNNING'}}},
            {'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}}
        ]
        mock_athena.get_query_results.return_value = {
            'ResultSet': {
                'Rows': [
                    {'Data': [{'VarCharValue': 'user_id'}, {'VarCharValue': 'name'}]},
                    {'Data': [{'VarCharValue': '201'}, {'VarCharValue': 'Charlie'}]}
                ]
            }
        }
        
        # Patch boto3.client to return our mock
        with patch('boto3.client', return_value=mock_athena):
            result = run_athena_query(
                query='SELECT * FROM user_data',
                database='analytics_db',
                output_location='s3://my-bucket/athena-results/'
            )
            
            # Validate results and method calls
            self.assertEqual(result[0]['name'], 'Charlie')
            mock_athena.start_query_execution.assert_called_once_with(
                QueryString='SELECT * FROM user_data',
                QueryExecutionContext={'Database': 'analytics_db'},
                ResultConfiguration={'OutputLocation': 's3://my-bucket/athena-results/'}
            )
            self.assertEqual(mock_athena.get_query_execution.call_count, 2)

Best Practices #

• Test Edge Cases: Mock scenarios like query timeouts, permission errors, or partial results to ensure your function handles them gracefully.
• Reuse Mock Setups: Use pytest fixtures to create reusable stubbed/mocked Athena clients across multiple tests.
• Validate API Calls: Use stubber.assert_no_pending_responses() (for Stubber) or mock.method.assert_called_with() (for MagicMock) to ensure your code makes the correct Athena API calls.
• Isolate Tests: Never rely on real AWS resources for unit tests—mocking ensures tests are fast, reliable, and don’t incur costs.

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