Athena tests refactoring.

Author: igorborgest

tests/_utils.py

@@ -4,6 +4,7 @@
 from decimal import Decimal
 
 import boto3
+import botocore.exceptions
 import pandas as pd
 
 import awswrangler as wr
@@ -470,3 +471,17 @@ def extract_cloudformation_outputs():
             for output in stack.get("Outputs"):
                 outputs[output.get("OutputKey")] = output.get("OutputValue")
     return outputs
+
+
+def list_workgroups():
+    client = boto3.client("athena")
+    attempt = 1
+    while True:
+        try:
+            return client.list_work_groups()
+        except botocore.exceptions.ClientError as ex:
+            if ex.response["Error"]["Code"] != "ThrottlingException":
+                raise ex
+            if attempt > 5:
+                raise ex
+            time.sleep(attempt + random.randrange(start=0, stop=3, step=1))

tests/conftest.py

@@ -5,7 +5,7 @@
 
 import awswrangler as wr
 
-from ._utils import extract_cloudformation_outputs, get_time_str_with_random_suffix, path_generator
+from ._utils import extract_cloudformation_outputs, get_time_str_with_random_suffix, list_workgroups, path_generator
 
 
 @pytest.fixture(scope="session")
@@ -66,7 +66,7 @@ def loggroup(cloudformation_outputs):
 def workgroup0(bucket):
     wkg_name = "aws_data_wrangler_0"
     client = boto3.client("athena")
-    wkgs = client.list_work_groups()
+    wkgs = list_workgroups()
     wkgs = [x["Name"] for x in wkgs["WorkGroups"]]
     if wkg_name not in wkgs:
         client.create_work_group(
@@ -87,7 +87,7 @@ def workgroup0(bucket):
 def workgroup1(bucket):
     wkg_name = "aws_data_wrangler_1"
     client = boto3.client("athena")
-    wkgs = client.list_work_groups()
+    wkgs = list_workgroups()
     wkgs = [x["Name"] for x in wkgs["WorkGroups"]]
     if wkg_name not in wkgs:
         client.create_work_group(
@@ -111,7 +111,7 @@ def workgroup1(bucket):
 def workgroup2(bucket, kms_key):
     wkg_name = "aws_data_wrangler_2"
     client = boto3.client("athena")
-    wkgs = client.list_work_groups()
+    wkgs = list_workgroups()
     wkgs = [x["Name"] for x in wkgs["WorkGroups"]]
     if wkg_name not in wkgs:
         client.create_work_group(
@@ -135,7 +135,7 @@ def workgroup2(bucket, kms_key):
 def workgroup3(bucket, kms_key):
     wkg_name = "aws_data_wrangler_3"
     client = boto3.client("athena")
-    wkgs = client.list_work_groups()
+    wkgs = list_workgroups()
     wkgs = [x["Name"] for x in wkgs["WorkGroups"]]
     if wkg_name not in wkgs:
         client.create_work_group(
@@ -199,6 +199,7 @@ def glue_table(glue_database):
     wr.catalog.delete_table_if_exists(database=glue_database, table=name)
     yield name
     wr.catalog.delete_table_if_exists(database=glue_database, table=name)
+    print(f"Table {glue_database}.{name} deleted.")
 
 
 @pytest.fixture(scope="function")