From 39da731e92ea11f45a944b5ba38be47461be2ad9 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 23 Jul 2026 13:56:55 -0700 Subject: [PATCH 1/2] 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 Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../microsoft/azure/sensors/compute.py | 10 ++++--- .../microsoft/azure/sensors/test_compute.py | 27 ++++++++++++++----- .../validate_operators_init_exemptions.txt | 1 - 3 files changed, 26 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 1a9c05185dcb6..6fc50516fb6e6 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 @@ def __init__( 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,12 @@ def __init__( self.deferrable = deferrable def poke(self, context: Context) -> bool: + # target_state is a template field; validate the rendered value here rather than in + # __init__, which only sees the un-rendered Jinja expression. + 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 47e4da37a8d46..809c3fe159510 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,27 @@ def test_init(self): 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): + # __init__ no longer validates the template field; the ValueError surfaces at poke, + # once target_state has been rendered. + 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): + # A templated target_state must construct: in __init__ it is still the un-rendered expression. + 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 01fd5bc56dbb7..af9658b934c99 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -73,7 +73,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 From 409164833107fe27863f99211e960068fa64334a Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 23 Jul 2026 23:20:00 -0700 Subject: [PATCH 2/2] Tighten AzureVirtualMachineStateSensor comments Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../src/airflow/providers/microsoft/azure/sensors/compute.py | 3 +-- .../azure/tests/unit/microsoft/azure/sensors/test_compute.py | 3 --- 2 files changed, 1 insertion(+), 5 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 6fc50516fb6e6..648b111432d78 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 @@ -68,8 +68,7 @@ def __init__( self.deferrable = deferrable def poke(self, context: Context) -> bool: - # target_state is a template field; validate the rendered value here rather than in - # __init__, which only sees the un-rendered Jinja expression. + # 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)}" 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 809c3fe159510..85c168ab89bde 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 @@ -44,8 +44,6 @@ def test_init(self): assert sensor.azure_conn_id == CONN_ID def test_invalid_target_state_rejected_at_poke(self): - # __init__ no longer validates the template field; the ValueError surfaces at poke, - # once target_state has been rendered. sensor = AzureVirtualMachineStateSensor( task_id="sense_vm", resource_group_name=RESOURCE_GROUP, @@ -56,7 +54,6 @@ def test_invalid_target_state_rejected_at_poke(self): sensor.poke(context=None) def test_templated_target_state_constructs(self): - # A templated target_state must construct: in __init__ it is still the un-rendered expression. sensor = AzureVirtualMachineStateSensor( task_id="sense_vm", resource_group_name=RESOURCE_GROUP,