amoghrajesh commented on code in PR #76918:
URL: https://github.com/apache/airflow/pull/76918#discussion_r4546388828
##########
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"},
+ ],
+ ids=["error", "cancelled", "not_found"],
+ )
+ def test_stored_run_not_reusable_submits_fresh(self, mock_trigger,
mock_wait, requests_mock, response):
+ self.task_store.get.return_value = STORED_RUN_ID
+ requests_mock.get(STORED_RUN_API_URL, **response)
+
+ assert self._operator().execute(self.context) == RUN_ID
+
+ mock_trigger.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(
+ ("stored_run_id", "expected_run_id"),
+ [(None, RUN_ID), (STORED_RUN_ID, STORED_RUN_ID)],
+ ids=["fresh", "reconnect"],
+ )
+ @patch.object(DbtCloudRunJobOperator, "defer", autospec=True)
+ @patch.object(
+ DbtCloudHook, "get_job_run_status", autospec=True,
return_value=DbtCloudJobRunStatus.RUNNING.value
+ )
+ def test_deferrable_defers_on_the_durable_run(
+ self, mock_status, mock_defer, mock_trigger, mock_wait, requests_mock,
stored_run_id, expected_run_id
+ ):
+ self.task_store.get.return_value = stored_run_id
+ requests_mock.get(
+ STORED_RUN_API_URL,
+ json={
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.RUNNING.value,
+ "href": STORED_RUN_URL,
+ }
+ },
+ )
+
+ self._operator(deferrable=True).execute(self.context)
+
+ assert mock_trigger.called is (stored_run_id is None)
+ mock_wait.assert_not_called()
+ assert mock_defer.call_args.kwargs["trigger"].run_id == expected_run_id
+ if stored_run_id is None:
Review Comment:
Split into two tests. The reconnect test now asserts `task_store.set` isn't
called.
--
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]
