amoghrajesh commented on code in PR #76918:
URL: https://github.com/apache/airflow/pull/76918#discussion_r4546377852
##########
providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py:
##########
@@ -1058,6 +1063,150 @@ def test_run_job_operator_link(
)
+STORED_RUN_ID = 9999
+STORED_RUN_URL = EXPECTED_JOB_RUN_OP_EXTRA_LINK.format(
+ account_id=DEFAULT_ACCOUNT_ID, project_id=PROJECT_ID, run_id=STORED_RUN_ID
+)
+STORED_RUN_API_URL =
f"https://cloud.getdbt.com/api/v2/accounts/{DEFAULT_ACCOUNT_ID}/runs/{STORED_RUN_ID}/";
+
+
[email protected](
+ not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
+)
[email protected](DbtCloudHook, "wait_for_job_run_status", autospec=True,
return_value=True)
[email protected](
+ DbtCloudHook,
+ "trigger_job_run",
+ autospec=True,
+ return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE),
+)
+class TestDbtCloudRunJobOperatorDurable:
+ def setup_method(self):
+ self.dag = DAG("test_dbt_cloud_job_run_op_durable", schedule=None,
start_date=DEFAULT_DATE)
+ self.ti = MagicMock(stats_tags={}, try_number=1)
+ self.task_store = MagicMock(spec_set=["get", "set"])
+ self.task_store.get.return_value = None
+ self.context = {"ti": self.ti, "task_state_store": self.task_store}
+
+ def _operator(self, **kwargs):
+ return DbtCloudRunJobOperator(
+ task_id=TASK_ID, dbt_cloud_conn_id=ACCOUNT_ID_CONN, job_id=JOB_ID,
dag=self.dag, **kwargs
+ )
+
+ def _assert_run_xcoms(self, run_id, run_url):
+ assert self.ti.xcom_push.call_args_list == [
+ call(key="job_run_url", value=run_url),
+ call(key="job_run_id", value=run_id),
+ ]
+
+ def test_fresh_submit_persists_run_id(self, mock_trigger, mock_wait):
+ assert self._operator().execute(self.context) == RUN_ID
+
+ mock_trigger.assert_called_once()
+ mock_wait.assert_called_once()
+ self.task_store.set.assert_called_once_with("dbt_cloud_job_run_id",
RUN_ID)
+ self._assert_run_xcoms(RUN_ID,
DEFAULT_ACCOUNT_JOB_RUN_RESPONSE["data"]["href"])
+
+ @pytest.mark.parametrize(
+ ("status", "waits"),
+ [
+ (DbtCloudJobRunStatus.QUEUED, True),
+ (DbtCloudJobRunStatus.STARTING, True),
+ (DbtCloudJobRunStatus.RUNNING, True),
+ (DbtCloudJobRunStatus.SUCCESS, False),
+ ],
+ )
+ def test_stored_run_is_reused(self, mock_trigger, mock_wait,
requests_mock, status, waits):
+ self.task_store.get.return_value = STORED_RUN_ID
+ requests_mock.get(
+ STORED_RUN_API_URL,
+ json={"data": {"id": STORED_RUN_ID, "status": status.value,
"href": STORED_RUN_URL}},
+ )
+ operator = self._operator()
+
+ assert operator.execute(self.context) == STORED_RUN_ID
+
+ mock_trigger.assert_not_called()
+ assert mock_wait.called is waits
+ self.task_store.set.assert_not_called()
+ assert operator.run_id == STORED_RUN_ID
+ self._assert_run_xcoms(STORED_RUN_ID, STORED_RUN_URL)
+
+ @pytest.mark.parametrize(
+ "response",
+ [
+ {
+ "json": {
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.ERROR.value,
+ "href": STORED_RUN_URL,
+ }
+ }
+ },
+ {
+ "json": {
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.CANCELLED.value,
+ "href": STORED_RUN_URL,
+ }
+ }
+ },
+ {"status_code": 404, "reason": "Not Found"},
Review Comment:
Added `test_stored_run_lookup_error_fails_without_triggering`. It's the only
test that fails if `== 404` is widened to `>= 400`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
