python-3.xamazon-web-servicespytestboto3moto

How to mock Athena query results values with Moto3 for a specific table?


I am using pytest and moto3 to test some code similar to this:

response = athena_client.start_query_execution(
        QueryString='SELECT * FROM xyz',
        QueryExecutionContext={'Database': myDb},
        ResultConfiguration={'OutputLocation': someLocation},
        WorkGroup=myWG
    )

execution_id = response['QueryExecutionId']

if response['QueryExecution']['Status']['State'] == 'SUCCEEDED':
    response = athena_client.get_query_results(
        QueryExecutionId=execution_id
    )

    results = response['ResultSet']['Rows']
    ...etc

In my test I need that the values from results = response['ResultSet']['Rows'] are controlled by the test. I am using some code like this:

backend = athena_backends[DEFAULT_ACCOUNT_ID]["us-east-1"]
    rows = [{"Data": [{"VarCharValue": "xyz"}]}, {"Data": [{"VarCharValue": ...}, etc]}]
    column_info = [
        {
            "CatalogName": "string",
            "SchemaName": "string",
            "TableName": "xyz",
            "Name": "string",
            "Label": "string",
            "Type": "string",
            "Precision": 123,
            "Scale": 123,
            "Nullable": "NOT_NULL",
            "CaseSensitive": True,
        }
    ]
    results = QueryResults(rows=rows, column_info=column_info)
    backend.query_results[NEEDED_QUERY_EXECUTION_ID] = results

but that is not working as I guess NEEDED_QUERY_EXECUTION_ID is not known before from the test. How can I control it?

UPDATE

Based on suggestion I tried to use:

results = QueryResults(rows=rows, column_info=column_info)
d = defaultdict(lambda: results.to_dict())
backend.query_results = d

to force a return of values, but it seems not working as from the moto3's models.AthenaBackend.get_query_results, I have this code:

    results = (
        self.query_results[exec_id]
        if exec_id in self.query_results
        else QueryResults(rows=[], column_info=[])
    )
    return results

which will fail as the if condition won't be satifsfied.


Solution

  • Extending the solution of the defaultdict, you could create a custom dictionary that contains all execution_ids, and always returns the same object:

    class QueryDict(dict):
        def __contains__(self, item):
            return True
         def __getitem__(self, item):
            rows = [{"Data": [{"VarCharValue": "xyz"}]}, {"Data": [{"VarCharValue": "..."}]}]
            column_info = [
                {
                    "CatalogName": "string",
                    "SchemaName": "string",
                    "TableName": "xyz",
                    "Name": "string",
                    "Label": "string",
                    "Type": "string",
                    "Precision": 123,
                    "Scale": 123,
                    "Nullable": "NOT_NULL",
                    "CaseSensitive": True,
                }
            ]
            return QueryResults(rows=rows, column_info=column_info)
    
    backend = athena_backends[DEFAULT_ACCOUNT_ID]["us-east-1"]
    
    backend.query_results = QueryDict()