fix (#123)

Co-authored-by: Maxim Mityutko <[email protected]>

brickflow_plugins/databricks/workflow_dependency_sensor.py

@@ -256,13 +256,19 @@ def execute(self):
  for task in run.tasks:
  if task.task_key == self.dependency_task_name:
  task_state = task.state.result_state
- self.log.info(
- f"Found the run_id '{run.run_id}' and '{self.dependency_task_name}' "
- f"task with state: {task_state.value}"
- )
- if task_state.value == "SUCCESS":
- self.log.info(f"Found a successful run: {run.run_id}")
- return
+ if task_state:
+ self.log.info(
+ f"Found the run_id '{run.run_id}' and '{self.dependency_task_name}' "
+ f"task with state: {task_state.value}"
+ )
+ if task_state.value == "SUCCESS":
+ self.log.info(f"Found a successful run: {run.run_id}")
+ return
+ else:
+ self.log.info(
+ f"Found the run_id '{run.run_id}' and '{self.dependency_task_name}' "
+ f"but the task has not started yet..."
+ )
 
  self.log.info("Didn't find a successful task run yet...")
 

tests/databricks_plugins/test_workflow_task_dependency_sensor.py

@@ -36,6 +36,11 @@ class TestWorkflowTaskDependencySensor:
  "result_state": "FAILED",
  },
  },
+ {
+ "run_id": 300,
+ "task_key": "baz",
+ "state": {},
+ },
  ],
  }
  ]
@@ -95,9 +100,30 @@ def test_sensor_failure(self, caplog, api):
 
  with pytest.raises(WorkflowDependencySensorTimeOutException):
  sensor.execute()
- assert "The job has timed out..." in caplog.text
- assert "Didn't find a successful task run yet..." in caplog.text
- assert (
- "Found the run_id '1' and 'bar' task with state: FAILED"
- in caplog.text
- )
+
+ assert (
+ "Found the run_id '1' and 'bar' task with state: FAILED"
+ in caplog.messages
+ )
+ assert "Didn't find a successful task run yet..." in caplog.messages
+
+ def test_sensor_no_state(self, caplog, api):
+ with api:
+ sensor = WorkflowTaskDependencySensor(
+ databricks_host=self.workspace_url,
+ databricks_token="token",
+ dependency_job_name="job",
+ dependency_task_name="baz",
+ delta=timedelta(seconds=1),
+ timeout_seconds=1,
+ poke_interval_seconds=1,
+ )
+
+ with pytest.raises(WorkflowDependencySensorTimeOutException):
+ sensor.execute()
+
+ assert (
+ "Found the run_id '1' and 'baz' but the task has not started yet..."
+ in caplog.messages
+ )
+ assert "Didn't find a successful task run yet..." in caplog.messages