renu-j commented on code in PR #61951:
URL: https://github.com/apache/airflow/pull/61951#discussion_r2905065722
##########
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py:
##########
@@ -805,11 +805,47 @@ def _retry_cluster_creation(self, hook: DataprocHook):
self.log.info("Cluster created.")
return Cluster.to_dict(cluster)
+ def _reconcile_cluster_state(self, hook: DataprocHook, cluster: Cluster)
-> Cluster:
+
+ if cluster.status.state == cluster.status.State.CREATING:
+ self.log.info("Cluster %s is in CREATING state.",
self.cluster_name)
+
+ cluster = self._wait_for_cluster_in_creating_state(hook)
+ self._handle_error_state(hook, cluster)
+ elif cluster.status.state == cluster.status.State.DELETING:
+ self.log.info("Cluster %s is in DELETING state.",
self.cluster_name)
+
+ self._wait_for_cluster_in_deleting_state(hook)
+
+ self.log.info("Attempting to re-create cluster: %s",
self.cluster_name)
+
+ operation = self._create_cluster(hook)
+ hook.wait_for_operation(
+ timeout=self.timeout,
+ result_retry=self.retry,
+ operation=operation,
+ )
+ cluster = self._get_cluster(hook)
+
Review Comment:
After recreating the cluster in the DELETING branch and
waiting for the operation to complete, should we verify
that the cluster transitions to RUNNING before returning?
##########
providers/google/tests/unit/google/cloud/operators/test_dataproc.py:
##########
@@ -819,7 +819,7 @@ def test_execute(self, mock_hook, to_dict_mock):
# Test whether xcom push occurs before create cluster is called
self.extra_links_manager_mock.assert_has_calls(expected_calls,
any_order=False)
- to_dict_mock.assert_called_once_with(mock_hook().wait_for_operation())
+
to_dict_mock.assert_called_once_with(mock_hook.return_value.get_cluster.return_value)
if AIRFLOW_V_3_0_PLUS:
Review Comment:
assert_called_once_with() is used twice with different arguments. Since this
assertion ensures the function is called exactly once, the second assertion
would contradict the first if multiple calls occur.
If the intention is to verify multiple calls with different arguments,
assert_has_calls() may be more appropriate.
--
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]
