diff --git a/tests/test_athena/test_athena_server_api.py b/tests/test_athena/test_athena_server_api.py index 46dc1b5f3..3955ee15d 100644 --- a/tests/test_athena/test_athena_server_api.py +++ b/tests/test_athena/test_athena_server_api.py @@ -1,5 +1,9 @@ +import os +from unittest import SkipTest + import boto3 import requests +from botocore.config import Config from moto import mock_aws, moto_proxy, settings @@ -21,6 +25,8 @@ DEFAULT_COLUMN_INFO = [ @mock_aws def test_set_athena_result(): + if settings.is_test_proxy_mode(): + raise SkipTest("Proxy requires more config, done in separate test") base_url = ( "localhost:5000" if settings.TEST_SERVER_MODE else "motoapi.amazonaws.com" ) @@ -28,20 +34,13 @@ def test_set_athena_result(): athena_result = { "results": [ { - "rows": [ - {"Data": [{"VarCharValue": "1"}]}, - ], + "rows": [{"Data": [{"VarCharValue": "1"}]}], "column_info": DEFAULT_COLUMN_INFO, } ] } - kwargs = {} - if settings.is_test_proxy_mode(): - add_proxy_details(kwargs) resp = requests.post( - f"http://{base_url}/moto-api/static/athena/query-results", - json=athena_result, - **kwargs, + f"http://{base_url}/moto-api/static/athena/query-results", json=athena_result ) assert resp.status_code == 201 @@ -59,6 +58,48 @@ def test_set_athena_result(): assert details["Rows"] == [] +def test_set_athena_result_using_proxy(): + if not settings.is_test_proxy_mode(): + raise SkipTest("Specifically testing the proxy here") + os.environ["AWS_ACCESS_KEY_ID"] = "testing" + os.environ["AWS_SECRET_ACCESS_KEY"] = "testing" + athena_result = { + "results": [ + { + "rows": [{"Data": [{"VarCharValue": "1"}]}], + "column_info": DEFAULT_COLUMN_INFO, + } + ] + } + kwargs = {"json": athena_result} + add_proxy_details(kwargs) + + # Delete any existing data + requests.post("http://motoapi.amazonaws.com/moto-api/reset", **kwargs) + # Store the QueryResults + resp = requests.post( + "http://motoapi.amazonaws.com/moto-api/static/athena/query-results", **kwargs + ) + assert resp.status_code == 201 + + proxy_config = Config(proxies={"http": "localhost:5005", "https": "localhost:5005"}) + client = boto3.client( + "athena", region_name="us-east-1", config=proxy_config, verify=False + ) + + details = client.get_query_results(QueryExecutionId="anyid")["ResultSet"] + assert details["Rows"] == athena_result["results"][0]["rows"] + assert details["ResultSetMetadata"]["ColumnInfo"] == DEFAULT_COLUMN_INFO + + # Operation should be idempotent + details = client.get_query_results(QueryExecutionId="anyid")["ResultSet"] + assert details["Rows"] == athena_result["results"][0]["rows"] + + # Different ID should still return different (default) results though + details = client.get_query_results(QueryExecutionId="otherid")["ResultSet"] + assert details["Rows"] == [] + + @mock_aws def test_set_multiple_athena_result(): base_url = (