Improving tests cleanup

Author: igorborgest

awswrangler/pandas.py

@@ -916,17 +916,19 @@ def data_to_s3(self,
     @staticmethod
     def _data_to_s3_dataset_writer(dataframe: pd.DataFrame,
                                    path: str,
-                                   partition_cols,
+                                   partition_cols: Optional[List[str]],
                                    preserve_index: bool,
                                    compression,
                                    session_primitives: "SessionPrimitives",
                                    file_format: str,
                                    cast_columns=None,
                                    extra_args: Optional[Dict[str, Optional[str]]] = None,
                                    isolated_dataframe: bool = False):
-        objects_paths = []
+        objects_paths: List[str] = []
         dataframe = Pandas._cast_pandas(dataframe=dataframe, cast_columns=cast_columns)
-        cast_columns_materialized = {c: t for c, t in cast_columns.items() if c not in partition_cols}
+        if partition_cols is None:
+            partition_cols = []
+        cast_columns_materialized: Dict[str, str] = {c: t for c, t in cast_columns.items() if c not in partition_cols}
         if not partition_cols:
             object_path = Pandas._data_to_s3_object_writer(dataframe=dataframe,
                                                            path=path,
@@ -1001,7 +1003,7 @@ def _data_to_s3_object_writer(dataframe: pd.DataFrame,
                                   compression: str,
                                   session_primitives: "SessionPrimitives",
                                   file_format: str,
-                                  cast_columns: Optional[List[str]] = None,
+                                  cast_columns: Optional[Dict[str, str]] = None,
                                   extra_args: Optional[Dict[str, Optional[str]]] = None,
                                   isolated_dataframe=False) -> str:
         fs = get_fs(session_primitives=session_primitives)

demo/getting_started.ipynb

@@ -1,8 +1,8 @@
 numpy~=1.18.1
 pandas~=0.25.3
 pyarrow~=0.15.1
-botocore~=1.13.47
-boto3~=1.10.47
+botocore~=1.13.48
+boto3~=1.10.48
 s3fs~=0.4.0
 tenacity~=6.0.0
 pg8000~=1.13.2

requirements.txt

@@ -6,6 +6,7 @@
 import boto3
 import pandas as pd
 
+import awswrangler as wr
 from awswrangler import Session, Athena
 from awswrangler.exceptions import QueryCancelled, QueryFailed
 
@@ -34,6 +35,10 @@ def database(cloudformation_outputs):
     else:
         raise Exception("You must deploy the test infrastructure using Cloudformation!")
     yield database
+    tables = wr.glue.tables(database=database)["Table"].tolist()
+    for t in tables:
+        print(f"Dropping: {database}.{t}...")
+        wr.glue.delete_table_if_exists(database=database, table=t)
 
 
 @pytest.fixture(scope="module")

testing/test_awswrangler/test_athena.py

@@ -43,6 +43,10 @@ def database(cloudformation_outputs):
     else:
         raise Exception("You must deploy the test infrastructure using SAM!")
     yield database
+    tables = wr.glue.tables(database=database)["Table"].tolist()
+    for t in tables:
+        print(f"Dropping: {database}.{t}...")
+        wr.glue.delete_table_if_exists(database=database, table=t)
 
 
 @pytest.fixture(scope="module")