From 9dbcb3911b2fd173a1e8d2584fa48a53c0474b93 Mon Sep 17 00:00:00 2001
From: Teddy <teddy.crepineau@gmail.com>
Date: Thu, 1 Sep 2022 13:47:00 +0200
Subject: [PATCH] Fix minor column data quality test bugs (#7111)

* Fixed test name issue + filtered out partition details for non BQ tables

* Exclude non BQ table from partition processing

* Fixed test + formating
---
 ingestion/src/metadata/orm_profiler/api/workflow.py   | 8 +++++++-
 ingestion/src/metadata/test_suite/api/workflow.py     | 8 +++++++-
 ingestion/src/metadata/test_suite/validations/core.py | 4 ++--
 ingestion/tests/unit/test_suite/test_validations.py   | 6 +++---
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/ingestion/src/metadata/orm_profiler/api/workflow.py b/ingestion/src/metadata/orm_profiler/api/workflow.py
index d0d5ca1bcf8..a57ab165aa5 100644
--- a/ingestion/src/metadata/orm_profiler/api/workflow.py
+++ b/ingestion/src/metadata/orm_profiler/api/workflow.py
@@ -35,7 +35,10 @@ from metadata.generated.schema.entity.data.table import (
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
 )
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseService,
+    DatabaseServiceType,
+)
 from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
@@ -207,6 +210,9 @@ class ProfilerWorkflow:
         Args:
             entity: table entity
         """
+        # Should remove this with https://github.com/open-metadata/OpenMetadata/issues/5458
+        if entity.serviceType != DatabaseServiceType.BigQuery:
+            return None
         entity_config: TableConfig = self.get_config_for_entity(entity)
         if entity_config:
             return entity_config.partitionConfig
diff --git a/ingestion/src/metadata/test_suite/api/workflow.py b/ingestion/src/metadata/test_suite/api/workflow.py
index 50ac915bc68..b1a8a17bb36 100644
--- a/ingestion/src/metadata/test_suite/api/workflow.py
+++ b/ingestion/src/metadata/test_suite/api/workflow.py
@@ -30,7 +30,10 @@ from metadata.generated.schema.entity.data.table import IntervalType, Table
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
 )
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseService,
+    DatabaseServiceType,
+)
 from metadata.generated.schema.metadataIngestion.testSuitePipeline import (
     TestSuitePipeline,
 )
@@ -195,6 +198,9 @@ class TestSuiteWorkflow:
         Args:
             entity: table entity
         """
+        # Should remove this with https://github.com/open-metadata/OpenMetadata/issues/5458
+        if entity.serviceType != DatabaseServiceType.BigQuery:
+            return None
         if entity.tablePartition:
             if entity.tablePartition.intervalType in {
                 IntervalType.TIME_UNIT,
diff --git a/ingestion/src/metadata/test_suite/validations/core.py b/ingestion/src/metadata/test_suite/validations/core.py
index dfc0fa0a866..9f4386ba9d8 100644
--- a/ingestion/src/metadata/test_suite/validations/core.py
+++ b/ingestion/src/metadata/test_suite/validations/core.py
@@ -133,11 +133,11 @@ validation_enum_registry.add("columnValueStdDevToBeBetween")(
 
 # # Column Session Tests
 validation_enum_registry.add("columnValuesToBeNotInSet")(column_values_not_in_set)
-validation_enum_registry.add("columnValuesToBeInSet")(column_values_in_set)
+validation_enum_registry.add("ColumnValuesToBeInSet")(column_values_in_set)
 validation_enum_registry.add("columnValuesToMatchRegex")(column_values_to_match_regex)
 validation_enum_registry.add("columnValuesToNotMatchRegex")(
     column_values_to_not_match_regex
 )
-validation_enum_registry.add("columnValuesMissingCountToBeEqual")(
+validation_enum_registry.add("columnValuesMissingCount")(
     column_values_missing_count_to_be_equal
 )
diff --git a/ingestion/tests/unit/test_suite/test_validations.py b/ingestion/tests/unit/test_suite/test_validations.py
index 3b63139fe14..496e0e85f53 100644
--- a/ingestion/tests/unit/test_suite/test_validations.py
+++ b/ingestion/tests/unit/test_suite/test_validations.py
@@ -270,7 +270,7 @@ class testSuiteValidation(unittest.TestCase):
             ],
         )
 
-        res = validation_enum_registry.registry["columnValuesToBeInSet"](
+        res = validation_enum_registry.registry["ColumnValuesToBeInSet"](
             test_case=test_case,
             execution_date=EXECUTION_DATE.timestamp(),
             runner=self.runner,
@@ -292,7 +292,7 @@ class testSuiteValidation(unittest.TestCase):
             ],
         )
 
-        res = validation_enum_registry.registry["columnValuesMissingCountToBeEqual"](
+        res = validation_enum_registry.registry["columnValuesMissingCount"](
             test_case=test_case,
             execution_date=EXECUTION_DATE.timestamp(),
             runner=self.runner,
@@ -313,7 +313,7 @@ class testSuiteValidation(unittest.TestCase):
             ],
         )
 
-        res = validation_enum_registry.registry["columnValuesMissingCountToBeEqual"](
+        res = validation_enum_registry.registry["columnValuesMissingCount"](
             test_case=test_case,
             execution_date=EXECUTION_DATE.timestamp(),
             runner=self.runner,