[airflow] branch main updated: Fix DataplexDataQualityJobStatusSensor and add unit tests (#33440)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 3157002eeb Fix DataplexDataQualityJobStatusSensor and add unit tests (#33440)
3157002eeb is described below

commit 3157002eeb8d5e2acf62a32ffe7a3647adc23105
Author: Beata Kossakowska <10...@users.noreply.github.com>
AuthorDate: Sat Aug 19 22:04:36 2023 +0200

    Fix DataplexDataQualityJobStatusSensor and add unit tests (#33440)
    
    Fix for the test in the Dataplex DQ.
    
    Add unit tests for cought and fixed DataplexDataQualityJobStatusSensor
    failure
    
    ---------
    
    Co-authored-by: Beata Kossakowska <bk...@google.com>
    Co-authored-by: Maksim Moiseenkov <ma...@epam.com>
---
 airflow/providers/google/cloud/sensors/dataplex.py |  7 ++--
 .../google/cloud/sensors/test_dataplex.py          | 43 ++++++++++++++++++++++
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/google/cloud/sensors/dataplex.py b/airflow/providers/google/cloud/sensors/dataplex.py
index 887c7e47c6..e1581fe1da 100644
--- a/airflow/providers/google/cloud/sensors/dataplex.py
+++ b/airflow/providers/google/cloud/sensors/dataplex.py
@@ -167,7 +167,7 @@ class DataplexDataQualityJobStatusSensor(BaseSensorOperator):
         impersonation_chain: str | Sequence[str] | None = None,
         fail_on_dq_failure: bool = False,
         result_timeout: float = 60.0 * 10,
-        start_sensor_time: float = time.monotonic(),
+        start_sensor_time: float | None = None,
         *args,
         **kwargs,
     ) -> None:
@@ -185,10 +185,9 @@ class DataplexDataQualityJobStatusSensor(BaseSensorOperator):
         self.result_timeout = result_timeout
         self.start_sensor_time = start_sensor_time
 
-    def execute(self, context: Context) -> None:
-        super().execute(context)
-
     def _duration(self):
+        if not self.start_sensor_time:
+            self.start_sensor_time = time.monotonic()
         return time.monotonic() - self.start_sensor_time
 
     def poke(self, context: Context) -> bool:
diff --git a/tests/providers/google/cloud/sensors/test_dataplex.py b/tests/providers/google/cloud/sensors/test_dataplex.py
index 3b59fd1ac6..c9ae358efa 100644
--- a/tests/providers/google/cloud/sensors/test_dataplex.py
+++ b/tests/providers/google/cloud/sensors/test_dataplex.py
@@ -23,6 +23,7 @@ from google.api_core.gapic_v1.method import DEFAULT
 from google.cloud.dataplex_v1.types import DataScanJob
 
 from airflow import AirflowException
+from airflow.providers.google.cloud.hooks.dataplex import AirflowDataQualityScanResultTimeoutException
 from airflow.providers.google.cloud.sensors.dataplex import (
     DataplexDataQualityJobStatusSensor,
     DataplexTaskStateSensor,
@@ -144,3 +145,45 @@ class TestDataplexDataQualityJobStatusSensor:
         )
 
         assert result
+
+    def test_start_sensor_time(self):
+        sensor = DataplexDataQualityJobStatusSensor(
+            task_id=TASK_ID,
+            project_id=PROJECT_ID,
+            job_id=TEST_JOB_ID,
+            data_scan_id=TEST_DATA_SCAN_ID,
+            region=REGION,
+            api_version=API_VERSION,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            timeout=TIMEOUT,
+        )
+
+        assert sensor.start_sensor_time is None
+
+        duration_1 = sensor._duration()
+        duration_2 = sensor._duration()
+
+        assert sensor.start_sensor_time
+        assert 0 < duration_1 < duration_2
+
+    @mock.patch.object(DataplexDataQualityJobStatusSensor, "_duration")
+    def test_start_sensor_time_timeout(self, mock_duration):
+        result_timeout = 100
+        mock_duration.return_value = result_timeout + 1
+
+        sensor = DataplexDataQualityJobStatusSensor(
+            task_id=TASK_ID,
+            project_id=PROJECT_ID,
+            job_id=TEST_JOB_ID,
+            data_scan_id=TEST_DATA_SCAN_ID,
+            region=REGION,
+            api_version=API_VERSION,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            timeout=TIMEOUT,
+            result_timeout=result_timeout,
+        )
+
+        with pytest.raises(AirflowDataQualityScanResultTimeoutException):
+            sensor.poke(context={})