Athena: Add additional proxy tests (#7401)

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": [
+                "rows": [{"Data": [{"VarCharValue": "1"}]}],
-                    {"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",
+        f"http://{base_url}/moto-api/static/athena/query-results", json=athena_result
-        json=athena_result,
-        **kwargs,
     )
     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 = (