This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d4622c770f2 Validate AzureVirtualMachineStateSensor target_state after
rendering (#70329)
d4622c770f2 is described below
commit d4622c770f2094d633987fbce9028bb49baaf91c
Author: Stefan Wang <[email protected]>
AuthorDate: Fri Jul 24 00:43:47 2026 -0700
Validate AzureVirtualMachineStateSensor target_state after rendering
(#70329)
Validate AzureVirtualMachineStateSensor target_state after rendering
target_state is a template field, so it is rendered after __init__ runs. The
constructor validated it against VALID_STATES and raised ValueError there,
so a
templated target_state (e.g. "{{ params.state }}") could never be built:
__init__
saw the un-rendered expression and rejected it. Move the check into poke(),
which
runs after rendering, and store target_state verbatim in the constructor.
related: #70296
---
.../providers/microsoft/azure/sensors/compute.py | 9 ++++----
.../unit/microsoft/azure/sensors/test_compute.py | 24 +++++++++++++++-------
.../ci/prek/validate_operators_init_exemptions.txt | 1 -
3 files changed, 22 insertions(+), 12 deletions(-)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/compute.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/compute.py
index 1a9c05185dc..648b111432d 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/compute.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/compute.py
@@ -60,10 +60,6 @@ class AzureVirtualMachineStateSensor(BaseSensorOperator):
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs,
) -> None:
- if target_state not in self.VALID_STATES:
- raise ValueError(
- f"Invalid target_state: {target_state}. Must be one of
{sorted(self.VALID_STATES)}"
- )
super().__init__(**kwargs)
self.resource_group_name = resource_group_name
self.vm_name = vm_name
@@ -72,6 +68,11 @@ class AzureVirtualMachineStateSensor(BaseSensorOperator):
self.deferrable = deferrable
def poke(self, context: Context) -> bool:
+ # target_state is a template field; validate the rendered value here,
not in __init__.
+ if self.target_state not in self.VALID_STATES:
+ raise ValueError(
+ f"Invalid target_state: {self.target_state}. Must be one of
{sorted(self.VALID_STATES)}"
+ )
hook = AzureComputeHook(azure_conn_id=self.azure_conn_id)
current_state = hook.get_power_state(self.resource_group_name,
self.vm_name)
self.log.info("VM %s power state: %s", self.vm_name, current_state)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_compute.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_compute.py
index 47e4da37a8d..85c168ab89b 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_compute.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_compute.py
@@ -43,14 +43,24 @@ class TestAzureVirtualMachineStateSensor:
assert sensor.target_state == "running"
assert sensor.azure_conn_id == CONN_ID
- def test_init_invalid_target_state(self):
+ def test_invalid_target_state_rejected_at_poke(self):
+ sensor = AzureVirtualMachineStateSensor(
+ task_id="sense_vm",
+ resource_group_name=RESOURCE_GROUP,
+ vm_name=VM_NAME,
+ target_state="invalid_state",
+ )
with pytest.raises(ValueError, match="Invalid target_state"):
- AzureVirtualMachineStateSensor(
- task_id="sense_vm",
- resource_group_name=RESOURCE_GROUP,
- vm_name=VM_NAME,
- target_state="invalid_state",
- )
+ sensor.poke(context=None)
+
+ def test_templated_target_state_constructs(self):
+ sensor = AzureVirtualMachineStateSensor(
+ task_id="sense_vm",
+ resource_group_name=RESOURCE_GROUP,
+ vm_name=VM_NAME,
+ target_state="{{ params.state }}",
+ )
+ assert sensor.target_state == "{{ params.state }}"
def test_template_fields(self):
sensor = AzureVirtualMachineStateSensor(
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index ae99a1b3369..669597d0a3d 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -67,7 +67,6 @@
providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py
providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py::GCSToGCSOperator
providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_local.py::GCSToLocalFilesystemOperator
providers/google/src/airflow/providers/google/marketing_platform/operators/campaign_manager.py::GoogleCampaignManagerDeleteReportOperator
-providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/compute.py::AzureVirtualMachineStateSensor
providers/microsoft/azure/src/airflow/providers/microsoft/azure/transfers/gcs_to_wasb.py::GCSToAzureBlobStorageOperator
providers/microsoft/azure/src/airflow/providers/microsoft/azure/transfers/oracle_to_azure_data_lake.py::OracleToAzureDataLakeOperator
providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py::PsrpOperator
