如何使用Moto3模拟指定Athena表的查询结果值？

针对你遇到的无法提前获取Query Execution ID、defaultdict不生效的问题，这里提供两种直接可行的解决方案：

方案1：提前捕获Execution ID并注入预设结果 #

业务代码中调用start_query_execution后会直接返回包含QueryExecutionId的响应，你可以在测试流程中先触发该调用拿到ID，再将预设结果注入Moto的backend，之后再执行后续业务逻辑。

示例代码：

def test_athena_query():
    # 初始化Athena客户端与backend
    athena_client = boto3.client("athena", region_name="us-east-1")
    backend = athena_backends[DEFAULT_ACCOUNT_ID]["us-east-1"]
    
    # 1. 先执行start_query_execution获取execution_id
    exec_response = athena_client.start_query_execution(
        QueryString='SELECT * FROM xyz',
        QueryExecutionContext={'Database': 'myDb'},
        ResultConfiguration={'OutputLocation': 's3://test-bucket'},
        WorkGroup='myWG'
    )
    execution_id = exec_response['QueryExecutionId']
    
    # 2. 构造并注入预设查询结果
    rows = [
        {"Data": [{"VarCharValue": "xyz"}]},
        {"Data": [{"VarCharValue": "test_data"}]}
    ]
    column_info = [
        {
            "CatalogName": "awsdatacatalog",
            "SchemaName": "myDb",
            "TableName": "xyz",
            "Name": "column1",
            "Label": "column1",
            "Type": "string",
            "Precision": 0,
            "Scale": 0,
            "Nullable": "NOT_NULL",
            "CaseSensitive": True,
        }
    ]
    preset_results = QueryResults(rows=rows, column_info=column_info)
    backend.query_results[execution_id] = preset_results
    
    # 3. 执行业务代码或直接验证结果
    if exec_response['QueryExecution']['Status']['State'] == 'SUCCEEDED':
        query_results = athena_client.get_query_results(QueryExecutionId=execution_id)
        assert query_results['ResultSet']['Rows'] == rows

如果业务代码是封装好的函数无法拆分，可以用Monkey Patch替换start_query_execution，返回预设ID的同时提前注入结果：

def test_athena_patched_execution():
    athena_client = boto3.client("athena", region_name="us-east-1")
    backend = athena_backends[DEFAULT_ACCOUNT_ID]["us-east-1"]
    
    # 预设ID与结果
    mock_exec_id = "test-exec-12345"
    preset_rows = [{"Data": [{"VarCharValue": "patched_result"}]}]
    preset_results = QueryResults(rows=preset_rows, column_info=[{"Name": "col1", "Type": "string"}])
    backend.query_results[mock_exec_id] = preset_results
    
    # 替换原方法
    original_start = athena_client.start_query_execution
    def patched_start(**kwargs):
        return {
            'QueryExecutionId': mock_exec_id,
            'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}
        }
    athena_client.start_query_execution = patched_start
    
    try:
        # 执行你的业务代码
        your_athena_business_logic()
        # 后续验证结果逻辑
    finally:
        # 恢复原方法
        athena_client.start_query_execution = original_start

方案2：直接Patch AthenaBackend的get_query_results方法 #

既然Moto原方法会判断exec_id是否存在于query_results中，你可以直接替换该方法，绕过判断逻辑，强制返回预设结果。

示例代码：

import types
from moto.athena.models import QueryResults

def test_athena_patch_backend_method():
    athena_client = boto3.client("athena", region_name="us-east-1")
    backend = athena_backends[DEFAULT_ACCOUNT_ID]["us-east-1"]
    
    # 构造预设结果
    preset_rows = [{"Data": [{"VarCharValue": "forced_result"}]}]
    preset_results = QueryResults(rows=preset_rows, column_info=[{"Name": "col1", "Type": "string"}])
    
    # 替换get_query_results方法
    original_get_results = backend.get_query_results
    def patched_get_results(self, exec_id):
        return preset_results
    
    backend.get_query_results = types.MethodType(patched_get_results, backend)
    
    try:
        # 执行业务代码，所有get_query_results调用都会返回预设结果
        your_athena_business_logic()
    finally:
        # 恢复原方法
        backend.get_query_results = original_get_results

补充说明 #

你之前用defaultdict无效的原因是：Moto的get_query_results方法中，exec_id in self.query_results的判断不会触发defaultdict的默认值生成（仅当直接访问键时才会触发），因此判断会返回False，最终返回空结果。上述两种方案均绕过了这个判断逻辑，或提前注入了合法的Execution ID与对应结果。

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