Increase coverage DQ

Author: jaidisido

awswrangler/data_quality/_create.py

@@ -295,7 +295,7 @@ def evaluate_ruleset(
     catalog_id: Optional[str] = None,
     connection_name: Optional[str] = None,
     additional_options: Optional[Dict[str, str]] = None,
-    additional_run_options: Optional[Dict[str, str]] = None,
+    additional_run_options: Optional[Dict[str, Union[str, bool]]] = None,
     client_token: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> pd.DataFrame:
@@ -324,7 +324,7 @@ def evaluate_ruleset(
         `pushDownPredicate`: to filter on partitions without having to list and read all the files in your dataset.
         `catalogPartitionPredicate`: to use server-side partition pruning using partition indexes in the
         Glue Data Catalog.
-    additional_run_options : Dict[str, str], optional
+    additional_run_options : Dict[str, Union[str, bool]], optional
         Additional run options. Supported keys:
         `CloudWatchMetricsEnabled`: whether to enable CloudWatch metrics.
         `ResultsS3Prefix`: prefix for Amazon S3 to store results.

awswrangler/data_quality/_utils.py

@@ -75,7 +75,7 @@ def _start_ruleset_evaluation_run(
     catalog_id: Optional[str] = None,
     connection_name: Optional[str] = None,
     additional_options: Optional[Dict[str, str]] = None,
-    additional_run_options: Optional[Dict[str, str]] = None,
+    additional_run_options: Optional[Dict[str, Union[str, bool]]] = None,
     client_token: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> str:

tests/test_athena.py

@@ -1045,7 +1045,6 @@ def test_bucketing_csv_saving(path, glue_database, glue_table, dtype):
     df3 = wr.athena.read_sql_query(query, database=glue_database, params=query_params, ctas_approach=False)
     scanned_bucketed = df3.query_metadata["Statistics"]["DataScannedInBytes"]
 
-    print(scanned_bucketed)
     assert df2.equals(df3)
     assert scanned_regular >= scanned_bucketed * saving_factor
 

tests/test_data_quality.py

@@ -69,14 +69,15 @@ def test_ruleset_dqdl(df, path, glue_database, glue_table, glue_ruleset, glue_da
     assert df_results["Result"].eq("PASS").all()
 
 
-def test_recommendation_ruleset(df, path, glue_database, glue_table, glue_ruleset, glue_data_quality_role):
+@pytest.mark.parametrize("name", [False, True])
+def test_recommendation_ruleset(df, path, name, glue_database, glue_table, glue_ruleset, glue_data_quality_role):
     df_recommended_ruleset = wr.data_quality.create_recommendation_ruleset(
+        name=f"{glue_ruleset}_recommended" if name else None,
         database=glue_database,
         table=glue_table,
         iam_role_arn=glue_data_quality_role,
         number_of_workers=2,
     )
-    print(glue_data_quality_role)
     df_rules = df_recommended_ruleset.append(
         {"rule_type": "ColumnValues", "parameter": '"c2"', "expression": "in [0, 1, 2]"}, ignore_index=True
     )
@@ -90,11 +91,12 @@ def test_recommendation_ruleset(df, path, glue_database, glue_table, glue_rulese
         name=glue_ruleset,
         iam_role_arn=glue_data_quality_role,
         number_of_workers=2,
+        additional_run_options={"CloudWatchMetricsEnabled": False},
     )
     assert df_results["Result"].eq("PASS").all()
 
 
-def test_ruleset_fail(df, path, glue_database, glue_table, glue_ruleset, glue_data_quality_role):
+def test_ruleset_fail(df, path, glue_database, glue_table, glue_ruleset, glue_data_quality_role, account_id):
     wr.data_quality.create_ruleset(
         name=glue_ruleset,
         database=glue_database,
@@ -105,6 +107,7 @@ def test_ruleset_fail(df, path, glue_database, glue_table, glue_ruleset, glue_da
         name=glue_ruleset,
         iam_role_arn=glue_data_quality_role,
         number_of_workers=2,
+        catalog_id=account_id,
     )
     assert df_results["Result"][0] == "FAIL"
 
@@ -135,6 +138,13 @@ def test_create_ruleset_already_exists(
     glue_table: str,
     glue_ruleset: str,
 ) -> None:
+    with pytest.raises(wr.exceptions.InvalidArgumentCombination):
+        wr.data_quality.create_ruleset(
+            name=glue_ruleset,
+            database=glue_database,
+            table=glue_table,
+        )
+
     wr.data_quality.create_ruleset(
         name=glue_ruleset,
         database=glue_database,
@@ -182,7 +192,7 @@ def test_update_ruleset(df: pd.DataFrame, glue_database: str, glue_table: str, g
     assert df_rules.equals(df_ruleset)
 
 
-def test_update_ruleset_does_not_exists(df: pd.DataFrame, glue_ruleset: str) -> None:
+def test_update_ruleset_exceptions(df: pd.DataFrame, glue_ruleset: str) -> None:
     df_rules = pd.DataFrame(
         {
             "rule_type": ["RowCount"],
@@ -198,6 +208,12 @@ def test_update_ruleset_does_not_exists(df: pd.DataFrame, glue_ruleset: str) ->
             df_rules=df_rules,
         )
 
+    with pytest.raises(wr.exceptions.InvalidArgumentValue):
+        wr.data_quality.update_ruleset(name=glue_ruleset, df_rules=df_rules, mode="append")
+
+    with pytest.raises(wr.exceptions.InvalidArgumentCombination):
+        wr.data_quality.update_ruleset(name=glue_ruleset)
+
 
 def test_upsert_ruleset(df: pd.DataFrame, glue_database: str, glue_table: str, glue_ruleset: str) -> None:
     df_rules = pd.DataFrame(
@@ -222,15 +238,13 @@ def test_upsert_ruleset(df: pd.DataFrame, glue_database: str, glue_table: str, g
         }
     )
 
-    new_glue_ruleset_name = f"{glue_ruleset} 2.0"
     wr.data_quality.update_ruleset(
         name=glue_ruleset,
-        updated_name=new_glue_ruleset_name,
         mode="upsert",
         df_rules=df_upsert,
     )
 
-    df_ruleset = wr.data_quality.get_ruleset(name=new_glue_ruleset_name)
+    df_ruleset = wr.data_quality.get_ruleset(name=glue_ruleset)
 
     assert df_ruleset.shape == (5, 3)
     row_count = df_ruleset.loc[df_ruleset["rule_type"] == "RowCount"]

tests/test_opensearch.py

@@ -146,21 +146,18 @@ def domain_endpoint_elasticsearch_7_10_fgac(cloudformation_outputs):
 
 def test_connection_opensearch_1_0(domain_endpoint_opensearch_1_0):
     client = wr.opensearch.connect(host=domain_endpoint_opensearch_1_0)
-    print(client.info())
     assert len(client.info()) > 0
 
 
 def test_connection_opensearch_1_0_https(domain_endpoint_opensearch_1_0):
     client = wr.opensearch.connect(host=f"https://{domain_endpoint_opensearch_1_0}")
-    print(client.info())
     assert len(client.info()) > 0
 
 
 def test_connection_elasticsearch_7_10_fgac(domain_endpoint_elasticsearch_7_10_fgac, opensearch_password):
     client = wr.opensearch.connect(
         host=domain_endpoint_elasticsearch_7_10_fgac, username="test", password=opensearch_password
     )
-    print(client.info())
     assert len(client.info()) > 0