Skip to content

Commit

Permalink
Fix Azure Batch Hook instantation (#33731)
Browse files Browse the repository at this point in the history
The Hook instantiation for Azure Batch has been done in the
constructor, which is wrong. This has been detected when #33716 added
example dag and it started to fail provider imports as connection
has beeen missing to instantiate it.

The hook instantiation is now moved to cached property.
  • Loading branch information
potiuk authored Aug 25, 2023
1 parent 3bb1c07 commit 38f2737
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 1 deletion.
6 changes: 5 additions & 1 deletion airflow/providers/microsoft/azure/operators/batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations

from functools import cached_property
from typing import TYPE_CHECKING, Any, Sequence

from azure.batch import models as batch_models
Expand Down Expand Up @@ -176,7 +177,10 @@ def __init__(
self.timeout = timeout
self.should_delete_job = should_delete_job
self.should_delete_pool = should_delete_pool
self.hook = self.get_hook()

@cached_property
def hook(self):
return self.get_hook()

def _check_inputs(self) -> Any:
if not self.os_family and not self.vm_publisher:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,7 @@ def setup_method(self, method, mock_batch, mock_hook):
self.batch_client = mock_batch.return_value
self.mock_instance = mock_hook.return_value
assert self.batch_client == self.operator.hook.connection
assert self.batch_client == self.operator2_pass.hook.connection

@mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
def test_execute_without_failures(self, wait_mock):
Expand Down

0 comments on commit 38f2737

Please sign in to comment.