Change test execution method to run_dag() from operator.execute()

astronomer · Dec 14, 2022 · 2009843 · 2009843
1 parent c5a2214
commit 2009843
Show file tree

Hide file tree

Showing 6 changed files with 191 additions and 159 deletions.
diff --git a/python-sdk/src/astro/sql/operators/data_validations/ColumnCheckOperator.py b/python-sdk/src/astro/sql/operators/data_validations/ColumnCheckOperator.py
@@ -3,7 +3,6 @@
 import pandas
 from airflow import AirflowException
 from airflow.decorators.base import get_unique_task_id
-from airflow.models.xcom_arg import XComArg
 from airflow.providers.common.sql.operators.sql import SQLColumnCheckOperator
 
 from astro.databases import create_database
@@ -61,17 +60,22 @@ def __init__(
         self.kwargs = kwargs
         self.df = None
 
+        dataset_qualified_name = ""
+        dataset_conn_id = ""
+
         if isinstance(dataset, BaseTable):
             db = create_database(conn_id=self.dataset.conn_id)  # type: ignore
             self.conn_id = self.dataset.conn_id
-            super().__init__(
-                table=db.get_table_qualified_name(table=self.dataset),
-                column_mapping=self.column_mapping,
-                partition_clause=self.partition_clause,
-                conn_id=dataset.conn_id,
-                database=dataset.metadata.database,
-                task_id=task_id if task_id is not None else get_unique_task_id("column_check"),
-            )
+            dataset_qualified_name = db.get_table_qualified_name(table=self.dataset)
+            dataset_conn_id = dataset.conn_id
+
+        super().__init__(
+            table=dataset_qualified_name,
+            column_mapping=self.column_mapping,
+            partition_clause=self.partition_clause,
+            conn_id=dataset_conn_id,
+            task_id=task_id if task_id is not None else get_unique_task_id("column_check"),
+        )
 
     def execute(self, context: "Context"):
         if isinstance(self.dataset, BaseTable):
@@ -198,7 +202,7 @@ def column_check(
     partition_clause: Optional[str] = None,
     task_id: Optional[str] = None,
     **kwargs,
-) -> XComArg:
+) -> ColumnCheckOperator:
     """
     Performs one or more of the templated checks in the column_checks dictionary.
     Checks are performed on a per-column basis specified by the column_mapping.
@@ -235,4 +239,4 @@ def column_check(
         partition_clause=partition_clause,
         kwargs=kwargs,
         task_id=task_id,
-    ).output
+    )
diff --git a/python-sdk/tests/data/data_validation.csv b/python-sdk/tests/data/data_validation.csv
@@ -1,4 +1,4 @@
 name,age,city,emp_id
-Dwight Schrute,30,,10
+Dwight Schrute,30.0,,10
 Michael Scott,,LA,1
 Jim Halpert,,California City,35