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 33c2363add0 Move Compute Engine validation out of constructors (#70454)
33c2363add0 is described below
commit 33c2363add03a35964047980ff688775a48bbe90
Author: daniel.jin <[email protected]>
AuthorDate: Sun Jul 26 04:34:18 2026 -0400
Move Compute Engine validation out of constructors (#70454)
---
.../providers/google/cloud/operators/compute.py | 209 ++---
.../unit/google/cloud/operators/test_compute.py | 899 ++++++++++++++-------
.../ci/prek/validate_operators_init_exemptions.txt | 10 -
3 files changed, 735 insertions(+), 383 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/compute.py
b/providers/google/src/airflow/providers/google/cloud/operators/compute.py
index 46d44821377..17be2ebc279 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/compute.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/compute.py
@@ -65,7 +65,6 @@ class ComputeEngineBaseOperator(GoogleCloudBaseOperator):
self.gcp_conn_id = gcp_conn_id
self.api_version = api_version
self.impersonation_chain = impersonation_chain
- self._validate_inputs()
super().__init__(**kwargs)
def _validate_inputs(self) -> None:
@@ -165,19 +164,12 @@ class
ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
self.body = body
self.zone = zone
self.request_id = request_id
- if "name" in body:
- resource_id = self.body["name"]
- self._field_validator = None # Optional[GcpBodyFieldValidator]
self.retry = retry
self.timeout = timeout
self.metadata = metadata
self.recreate_if_machine_type_different =
recreate_if_machine_type_different
+ self.validate_body = validate_body
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
super().__init__(
resource_id=resource_id,
zone=zone,
@@ -208,8 +200,11 @@ class
ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
)
def _validate_all_body_fields(self) -> None:
- if self._field_validator:
- self._field_validator.validate(self.body)
+ if self.validate_body:
+ GcpBodyFieldValidator(
+ GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
+ api_version=self.api_version,
+ ).validate(self.body)
def _extract_machine_type(self, value: str | None) -> str | None:
if not value:
@@ -237,7 +232,7 @@ class
ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
def _create_instance(self, hook: ComputeEngineHook, context: Context) ->
dict:
"""Create the instance using the current body and return the created
instance as dict."""
- self._field_sanitizer.sanitize(self.body)
+
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE).sanitize(self.body)
self.log.info("Creating Instance with specified body: %s", self.body)
@@ -273,13 +268,18 @@ class
ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
If machine type drift is detected and
``recreate_if_machine_type_different=True``,
the existing instance is deleted and recreated using the requested
body.
"""
+ if "name" in self.body:
+ self.resource_id = self.body["name"]
+
+ self._validate_inputs()
+ self._validate_all_body_fields()
+ self.check_body_fields()
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
impersonation_chain=self.impersonation_chain,
)
- self._validate_all_body_fields()
- self.check_body_fields()
try:
existing_instance = hook.get_instance(
@@ -411,19 +411,12 @@ class
ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
self.source_instance_template = source_instance_template
self.body = body
self.zone = zone
- if "name" in body:
- resource_id = self.body["name"]
self.request_id = request_id
- self._field_validator = None # Optional[GcpBodyFieldValidator]
self.retry = retry
self.timeout = timeout
self.metadata = metadata
+ self.validate_body = validate_body
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
super().__init__(
resource_id=resource_id,
zone=zone,
@@ -435,8 +428,11 @@ class
ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
)
def _validate_all_body_fields(self) -> None:
- if self._field_validator:
- self._field_validator.validate(self.body)
+ if self.validate_body:
+ GcpBodyFieldValidator(
+ GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
+ api_version=self.api_version,
+ ).validate(self.body)
def _validate_inputs(self) -> None:
super()._validate_inputs()
@@ -447,12 +443,19 @@ class
ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
)
def execute(self, context: Context) -> dict:
+ if "name" in self.body:
+ self.resource_id = self.body["name"]
+
+ self._validate_inputs()
+ self._validate_all_body_fields()
+
+ field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
impersonation_chain=self.impersonation_chain,
)
- self._validate_all_body_fields()
try:
# Idempotence check (sort of) - we want to check if the new
Instance
# is already created and if is, then we assume it was created - we
do
@@ -475,7 +478,7 @@ class
ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
project_id=self.project_id or hook.project_id,
)
return Instance.to_dict(existing_instance)
- self._field_sanitizer.sanitize(self.body)
+ field_sanitizer.sanitize(self.body)
self.log.info("Creating Instance with specified body: %s", self.body)
hook.insert_instance(
body=self.body,
@@ -561,16 +564,10 @@ class
ComputeEngineDeleteInstanceOperator(ComputeEngineBaseOperator):
self.zone = zone
self.request_id = request_id
self.resource_id = resource_id
- self._field_validator = None # Optional[GcpBodyFieldValidator]
self.retry = retry
self.timeout = timeout
self.metadata = metadata
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
super().__init__(
project_id=project_id,
zone=zone,
@@ -587,6 +584,8 @@ class
ComputeEngineDeleteInstanceOperator(ComputeEngineBaseOperator):
raise AirflowException("The required parameter 'resource_id' is
missing. ")
def execute(self, context: Context) -> None:
+ self._validate_inputs()
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
@@ -660,6 +659,7 @@ class
ComputeEngineStartInstanceOperator(ComputeEngineBaseOperator):
raise AirflowException("The required parameter 'resource_id' is
missing. ")
def execute(self, context: Context) -> None:
+ self._validate_inputs()
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
@@ -718,6 +718,7 @@ class
ComputeEngineStopInstanceOperator(ComputeEngineBaseOperator):
raise AirflowException("The required parameter 'resource_id' is
missing. ")
def execute(self, context: Context) -> None:
+ self._validate_inputs()
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
@@ -794,11 +795,7 @@ class
ComputeEngineSetMachineTypeOperator(ComputeEngineBaseOperator):
**kwargs,
) -> None:
self.body = body
- self._field_validator: GcpBodyFieldValidator | None = None
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- SET_MACHINE_TYPE_VALIDATION_SPECIFICATION,
api_version=api_version
- )
+ self.validate_body = validate_body
super().__init__(
project_id=project_id,
zone=zone,
@@ -810,8 +807,11 @@ class
ComputeEngineSetMachineTypeOperator(ComputeEngineBaseOperator):
)
def _validate_all_body_fields(self) -> None:
- if self._field_validator:
- self._field_validator.validate(self.body)
+ if self.validate_body:
+ GcpBodyFieldValidator(
+ SET_MACHINE_TYPE_VALIDATION_SPECIFICATION,
+ api_version=self.api_version,
+ ).validate(self.body)
def _validate_inputs(self) -> None:
super()._validate_inputs()
@@ -819,12 +819,14 @@ class
ComputeEngineSetMachineTypeOperator(ComputeEngineBaseOperator):
raise AirflowException("The required parameter 'resource_id' is
missing. ")
def execute(self, context: Context) -> None:
+ self._validate_inputs()
+ self._validate_all_body_fields()
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
impersonation_chain=self.impersonation_chain,
)
- self._validate_all_body_fields()
ComputeInstanceDetailsLink.persist(
context=context,
project_id=self.project_id or hook.project_id,
@@ -958,18 +960,11 @@ class
ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
) -> None:
self.body = body
self.request_id = request_id
- if "name" in body:
- resource_id = self.body["name"]
- self._field_validator = None # Optional[GcpBodyFieldValidator]
self.retry = retry
self.timeout = timeout
self.metadata = metadata
+ self.validate_body = validate_body
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
super().__init__(
project_id=project_id,
zone="global",
@@ -992,8 +987,11 @@ class
ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
)
def _validate_all_body_fields(self) -> None:
- if self._field_validator:
- self._field_validator.validate(self.body)
+ if self.validate_body:
+ GcpBodyFieldValidator(
+ GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
+ api_version=self.api_version,
+ ).validate(self.body)
def _validate_inputs(self) -> None:
super()._validate_inputs()
@@ -1004,14 +1002,21 @@ class
ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
)
def execute(self, context: Context) -> dict:
+ if "name" in self.body:
+ self.resource_id = self.body["name"]
+
+ self._validate_inputs()
+ self._validate_all_body_fields()
+ self.check_body_fields()
+
+ field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
+ field_sanitizer.sanitize(self.body)
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
impersonation_chain=self.impersonation_chain,
)
- self._validate_all_body_fields()
- self.check_body_fields()
- self._field_sanitizer.sanitize(self.body)
try:
# Idempotence check (sort of) - we want to check if the new
Template
# is already created and if is, then we assume it was created by
previous run
@@ -1035,7 +1040,6 @@ class
ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
project_id=self.project_id or hook.project_id,
)
return InstanceTemplate.to_dict(existing_template)
- self._field_sanitizer.sanitize(self.body)
self.log.info("Creating Instance Template with specified body: %s",
self.body)
hook.insert_instance_template(
body=self.body,
@@ -1114,16 +1118,10 @@ class
ComputeEngineDeleteInstanceTemplateOperator(ComputeEngineBaseOperator):
) -> None:
self.request_id = request_id
self.resource_id = resource_id
- self._field_validator = None # Optional[GcpBodyFieldValidator]
self.retry = retry
self.timeout = timeout
self.metadata = metadata
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
super().__init__(
project_id=project_id,
zone="global",
@@ -1140,6 +1138,8 @@ class
ComputeEngineDeleteInstanceTemplateOperator(ComputeEngineBaseOperator):
raise AirflowException("The required parameter 'resource_id' is
missing.")
def execute(self, context: Context) -> None:
+ self._validate_inputs()
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
@@ -1232,17 +1232,7 @@ class
ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
) -> None:
self.body_patch = body_patch
self.request_id = request_id
- self._field_validator = None # GcpBodyFieldValidator | None
- if "name" not in self.body_patch:
- raise AirflowException(
- f"The body '{body_patch}' should contain at least name for the
new operator "
- f"in the 'name' field"
- )
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
+ self.validate_body = validate_body
super().__init__(
project_id=project_id,
zone="global",
@@ -1253,9 +1243,19 @@ class
ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
**kwargs,
)
+ def _validate_body_patch(self) -> None:
+ if "name" not in self.body_patch:
+ raise AirflowException(
+ f"The body '{self.body_patch}' should contain at least name
for the new operator "
+ "in the 'name' field"
+ )
+
def _validate_all_body_fields(self) -> None:
- if self._field_validator:
- self._field_validator.validate(self.body_patch)
+ if self.validate_body:
+ GcpBodyFieldValidator(
+ GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
+ api_version=self.api_version,
+ ).validate(self.body_patch)
def _validate_inputs(self) -> None:
super()._validate_inputs()
@@ -1263,6 +1263,12 @@ class
ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
raise AirflowException("The required parameter 'resource_id' is
missing.")
def execute(self, context: Context) -> dict:
+ self._validate_inputs()
+ self._validate_body_patch()
+ self._validate_all_body_fields()
+
+ field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
@@ -1305,7 +1311,7 @@ class
ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
)
)
new_body = deepcopy(old_body)
- self._field_sanitizer.sanitize(new_body)
+ field_sanitizer.sanitize(new_body)
new_body = merge(new_body, self.body_patch)
self.log.info("Calling insert instance template with updated body:
%s", new_body)
hook.insert_instance_template(body=new_body,
request_id=self.request_id, project_id=self.project_id)
@@ -1393,12 +1399,6 @@ class
ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
self.request_id = request_id
self.update_policy = update_policy
self._change_performed = False
- if api_version == "v1":
- raise AirflowException(
- "Api version v1 does not have update/patch "
- "operations for Instance Group Managers. Use beta"
- " api version or above"
- )
super().__init__(
project_id=project_id,
zone=zone,
@@ -1409,6 +1409,14 @@ class
ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
**kwargs,
)
+ def _validate_api_version(self) -> None:
+ if self.api_version == "v1":
+ raise AirflowException(
+ "Api version v1 does not have update/patch "
+ "operations for Instance Group Managers. Use beta "
+ "api version or above"
+ )
+
def _validate_inputs(self) -> None:
super()._validate_inputs()
if not self.resource_id:
@@ -1420,6 +1428,9 @@ class
ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
self._change_performed = True
def execute(self, context: Context) -> bool | None:
+ self._validate_inputs()
+ self._validate_api_version()
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
@@ -1531,17 +1542,10 @@ class
ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
) -> None:
self.body = body
self.request_id = request_id
- if "name" in body:
- resource_id = self.body["name"]
- self._field_validator = None # Optional[GcpBodyFieldValidator]
self.retry = retry
self.timeout = timeout
self.metadata = metadata
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
+ self.validate_body = validate_body
super().__init__(
project_id=project_id,
zone=zone,
@@ -1564,8 +1568,11 @@ class
ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
)
def _validate_all_body_fields(self) -> None:
- if self._field_validator:
- self._field_validator.validate(self.body)
+ if self.validate_body:
+ GcpBodyFieldValidator(
+ GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
+ api_version=self.api_version,
+ ).validate(self.body)
def _validate_inputs(self) -> None:
super()._validate_inputs()
@@ -1576,13 +1583,20 @@ class
ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
)
def execute(self, context: Context) -> dict:
+ if "name" in self.body:
+ self.resource_id = self.body["name"]
+
+ self._validate_inputs()
+ self._validate_all_body_fields()
+ self.check_body_fields()
+
+ field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
impersonation_chain=self.impersonation_chain,
)
- self._validate_all_body_fields()
- self.check_body_fields()
try:
# Idempotence check (sort of) - we want to check if the new
Instance Group Manager
# is already created and if isn't, we create new one
@@ -1603,7 +1617,7 @@ class
ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
project_id=self.project_id or hook.project_id,
)
return
InstanceGroupManager.to_dict(existing_instance_group_manager)
- self._field_sanitizer.sanitize(self.body)
+ field_sanitizer.sanitize(self.body)
self.log.info("Creating Instance Group Manager with specified body:
%s", self.body)
hook.insert_instance_group_manager(
body=self.body,
@@ -1687,15 +1701,10 @@ class
ComputeEngineDeleteInstanceGroupManagerOperator(ComputeEngineBaseOperator)
self.zone = zone
self.request_id = request_id
self.resource_id = resource_id
- self._field_validator = None # Optional[GcpBodyFieldValidator]
self.retry = retry
self.timeout = timeout
self.metadata = metadata
- if validate_body:
- self._field_validator = GcpBodyFieldValidator(
- GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version
- )
- self._field_sanitizer =
GcpBodyFieldSanitizer(GCE_INSTANCE_FIELDS_TO_SANITIZE)
+
super().__init__(
project_id=project_id,
zone=zone,
@@ -1712,6 +1721,8 @@ class
ComputeEngineDeleteInstanceGroupManagerOperator(ComputeEngineBaseOperator)
raise AirflowException("The required parameter 'resource_id' is
missing. ")
def execute(self, context: Context):
+ self._validate_inputs()
+
hook = ComputeEngineHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
diff --git a/providers/google/tests/unit/google/cloud/operators/test_compute.py
b/providers/google/tests/unit/google/cloud/operators/test_compute.py
index 5d8bc187b07..2ca63771275 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_compute.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_compute.py
@@ -147,16 +147,24 @@ class TestGceInstanceInsert:
request_id=None,
)
- def test_insert_instance_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineInsertInstanceOperator(
- project_id="",
- body=GCE_INSTANCE_BODY_API_CALL,
- zone=GCE_ZONE,
- task_id=TASK_ID,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_insert_instance_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineInsertInstanceOperator(
+ project_id="",
+ body=GCE_INSTANCE_BODY_API_CALL,
+ zone=GCE_ZONE,
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_insert_instance_should_not_throw_ex_when_project_id_none(self,
mock_hook):
@@ -190,68 +198,96 @@ class TestGceInstanceInsert:
project_id=None,
)
- def test_insert_instance_should_throw_ex_when_missing_zone(self):
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_insert_instance_should_throw_ex_when_missing_zone(self,
mock_hook):
+ op = ComputeEngineInsertInstanceOperator(
+ resource_id=GCE_RESOURCE_ID,
+ body=GCE_INSTANCE_BODY_API_CALL,
+ zone="",
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
with pytest.raises(AirflowException, match=r"The required parameter
'zone' is missing"):
- ComputeEngineInsertInstanceOperator(
- resource_id=GCE_RESOURCE_ID,
- body=GCE_INSTANCE_BODY_API_CALL,
- zone="",
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
- def test_insert_instance_should_throw_ex_when_missing_resource_id(self):
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_insert_instance_should_throw_ex_when_missing_resource_id(self,
mock_hook):
+ op = ComputeEngineInsertInstanceOperator(
+ project_id=GCP_PROJECT_ID,
+ zone=GCE_ZONE,
+ body=GCE_INSTANCE_BODY_WITHOUT_NAME_API_CALL,
+ task_id=TASK_ID,
+ resource_id="",
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
with pytest.raises(
AirflowException,
match=r"The required parameters 'resource_id' and "
r"body\['name'\] are missing\. Please, provide "
r"at least one of them",
):
- ComputeEngineInsertInstanceOperator(
- project_id=GCP_PROJECT_ID,
- zone=GCE_ZONE,
- body=GCE_INSTANCE_BODY_WITHOUT_NAME_API_CALL,
- task_id=TASK_ID,
- resource_id="",
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ op.execute(context=mock.MagicMock())
+ mock_hook.assert_not_called()
+
+ @pytest.mark.db_test
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
- def test_insert_instance_should_not_throw_ex_when_name_is_templated(self,
mock_hook):
+ def test_insert_instance_should_not_throw_ex_when_name_is_templated(
+ self,
+ mock_hook,
+ create_task_instance_of_operator,
+ ):
+ dag_id = "templated-instance-name"
+
get_instance_obj_mock = mock.MagicMock()
get_instance_obj_mock.__class__ = Instance
mock_hook.return_value.get_instance.side_effect = [
NotFound("Error message"),
get_instance_obj_mock,
]
+
body_with_templated_name = deepcopy(GCE_INSTANCE_BODY_API_CALL)
- body_with_templated_name["name"] = "{{ logical_date }}"
- op = ComputeEngineInsertInstanceOperator(
+ body_with_templated_name["name"] = "{{ dag.dag_id }}"
+
+ ti = create_task_instance_of_operator(
+ ComputeEngineInsertInstanceOperator,
+ dag_id=dag_id,
project_id=GCP_PROJECT_ID,
- resource_id=GCE_RESOURCE_ID,
body=body_with_templated_name,
zone=GCE_ZONE,
task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
- op.execute(context=mock.MagicMock())
+
+ rendered = ti.render_templates()
+
+ assert rendered.body["name"] == dag_id
+ assert rendered.resource_id is None
+
+ rendered.execute(context=mock.MagicMock())
+
+ assert rendered.resource_id == dag_id
+
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
+ expected_body = deepcopy(GCE_INSTANCE_BODY_API_CALL)
+ expected_body["name"] = dag_id
+
mock_hook.return_value.insert_instance.assert_called_once_with(
project_id=GCP_PROJECT_ID,
- body=body_with_templated_name,
+ body=expected_body,
zone=GCE_ZONE,
request_id=None,
)
@@ -358,7 +394,13 @@ class TestGceInstanceInsertFromTemplate:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
+ assert op.resource_id is None
+
op.execute(context=mock.MagicMock())
+
+ assert op.resource_id == GCP_INSTANCE_BODY_FROM_TEMPLATE["name"]
+
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
@@ -372,17 +414,25 @@ class TestGceInstanceInsertFromTemplate:
request_id=None,
)
- def
test_insert_instance_from_template_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineInsertInstanceFromTemplateOperator(
- project_id="",
- source_instance_template=SOURCE_INSTANCE_TEMPLATE,
- body=GCP_INSTANCE_BODY_FROM_TEMPLATE,
- zone=GCE_ZONE,
- task_id=TASK_ID,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def
test_insert_instance_from_template_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineInsertInstanceFromTemplateOperator(
+ project_id="",
+ source_instance_template=SOURCE_INSTANCE_TEMPLATE,
+ body=GCP_INSTANCE_BODY_FROM_TEMPLATE,
+ zone=GCE_ZONE,
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def
test_insert_instance_from_template_should_not_throw_ex_when_project_id_none(self,
mock_hook):
@@ -417,17 +467,25 @@ class TestGceInstanceInsertFromTemplate:
request_id=None,
)
- def
test_insert_instance_from_template_should_throw_ex_when_missing_zone(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'zone' is missing"):
- ComputeEngineInsertInstanceFromTemplateOperator(
- project_id=GCP_PROJECT_ID,
- zone="",
- source_instance_template=SOURCE_INSTANCE_TEMPLATE,
- body=GCP_INSTANCE_BODY_FROM_TEMPLATE,
- task_id=TASK_ID,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def
test_insert_instance_from_template_should_throw_ex_when_missing_zone(self,
mock_hook):
+ op = ComputeEngineInsertInstanceFromTemplateOperator(
+ project_id=GCP_PROJECT_ID,
+ zone="",
+ source_instance_template=SOURCE_INSTANCE_TEMPLATE,
+ body=GCP_INSTANCE_BODY_FROM_TEMPLATE,
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'zone' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
def
test_insert_instance_from_template_should_throw_ex_when_missing_source_instance_template(self):
with pytest.raises(
@@ -453,17 +511,26 @@ class TestGceInstanceInsertFromTemplate:
impersonation_chain=IMPERSONATION_CHAIN,
)
+ @pytest.mark.db_test
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
- def
test_insert_instance_from_template_should_not_throw_ex_when_name_is_templated(self,
mock_hook):
+ def
test_insert_instance_from_template_should_not_throw_ex_when_name_is_templated(
+ self, mock_hook, create_task_instance_of_operator
+ ):
+ dag_id = "templated-instance-name"
+
get_instance_obj_mock = mock.MagicMock()
get_instance_obj_mock.__class__ = Instance
mock_hook.return_value.get_instance.side_effect = [
NotFound("Error message"),
get_instance_obj_mock,
]
+
body_with_templated_name = deepcopy(GCP_INSTANCE_BODY_FROM_TEMPLATE)
- body_with_templated_name["name"] = "{{ execution_date }}"
- op = ComputeEngineInsertInstanceFromTemplateOperator(
+ body_with_templated_name["name"] = "{{ dag.dag_id }}"
+
+ ti = create_task_instance_of_operator(
+ ComputeEngineInsertInstanceFromTemplateOperator,
+ dag_id=dag_id,
project_id=GCP_PROJECT_ID,
source_instance_template=SOURCE_INSTANCE_TEMPLATE,
body=body_with_templated_name,
@@ -475,7 +542,16 @@ class TestGceInstanceInsertFromTemplate:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
- op.execute(context=mock.MagicMock())
+
+ rendered = ti.render_templates()
+
+ assert rendered.body["name"] == dag_id
+ assert rendered.resource_id is None
+
+ rendered.execute(context=mock.MagicMock())
+
+ assert rendered.resource_id == dag_id
+
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
@@ -484,7 +560,10 @@ class TestGceInstanceInsertFromTemplate:
mock_hook.return_value.insert_instance.assert_called_once_with(
project_id=GCP_PROJECT_ID,
zone=GCE_ZONE,
- body=body_with_templated_name,
+ body={
+ **GCP_INSTANCE_BODY_FROM_TEMPLATE,
+ "name": dag_id,
+ },
source_instance_template=SOURCE_INSTANCE_TEMPLATE,
request_id=None,
)
@@ -516,37 +595,48 @@ class TestGceInstanceDelete:
zone=GCE_ZONE,
)
- def test_delete_instance_should_throw_ex_when_missing_zone(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'zone' is missing"):
- ComputeEngineDeleteInstanceOperator(
- resource_id=GCE_RESOURCE_ID,
- zone="",
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_delete_instance_should_throw_ex_when_missing_zone(self,
mock_hook):
+ op = ComputeEngineDeleteInstanceOperator(
+ project_id=GCP_PROJECT_ID,
+ zone="",
+ resource_id=GCE_RESOURCE_ID,
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
- def test_delete_instance_should_throw_ex_when_missing_resource_id(self):
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'zone' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
+
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_delete_instance_should_throw_ex_when_missing_resource_id(self,
mock_hook):
+ op = ComputeEngineDeleteInstanceOperator(
+ resource_id="",
+ zone=GCE_ZONE,
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
with pytest.raises(AirflowException, match=r"The required parameter
'resource_id' is missing"):
- ComputeEngineDeleteInstanceOperator(
- resource_id="",
- zone=GCE_ZONE,
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
class TestGceInstanceStart:
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_start_instance_should_execute_successfully(self, mock_hook):
mock_hook.return_value.start_instance.return_value = True
+
op = ComputeEngineStartInstanceOperator(
project_id=GCP_PROJECT_ID,
zone=GCE_ZONE,
@@ -555,6 +645,7 @@ class TestGceInstanceStart:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
context = mock.MagicMock()
op.execute(context=context)
@@ -564,14 +655,16 @@ class TestGceInstanceStart:
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.start_instance.assert_called_once_with(
- zone=GCE_ZONE, resource_id=GCE_RESOURCE_ID,
project_id=GCP_PROJECT_ID
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ project_id=GCP_PROJECT_ID,
)
# Setting all the operator's input parameters as template dag_ids
# (could be anything else) just to test if the templating works for all
fields
@pytest.mark.db_test
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
- def test_start_instance_with_templates(self, _,
create_task_instance_of_operator, session):
+ def test_start_instance_with_templates(self, _,
create_task_instance_of_operator):
dag_id = "test_instance_start_with_templates"
ti = create_task_instance_of_operator(
ComputeEngineStartInstanceOperator,
@@ -590,19 +683,27 @@ class TestGceInstanceStart:
assert dag_id == rendered.gcp_conn_id
assert dag_id == rendered.api_version
- def test_instance_start_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineStartInstanceOperator(
- project_id="",
- zone=GCE_ZONE,
- resource_id=GCE_RESOURCE_ID,
- task_id=TASK_ID,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_instance_start_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineStartInstanceOperator(
+ project_id="",
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
- def test_instance_start_should_not_throw_ex_when_project_id_none(self, _):
+ def test_instance_start_should_not_throw_ex_when_project_id_none(self,
mock_hook):
op = ComputeEngineStartInstanceOperator(
zone=GCE_ZONE,
resource_id=GCE_RESOURCE_ID,
@@ -613,27 +714,54 @@ class TestGceInstanceStart:
context = mock.MagicMock()
op.execute(context=context)
- def test_instance_start_should_throw_ex_when_missing_zone(self):
- with pytest.raises(AirflowException, match=r"he required parameter
'zone' is missing"):
- ComputeEngineStartInstanceOperator(
- project_id=GCP_PROJECT_ID,
- zone="",
- resource_id=GCE_RESOURCE_ID,
- task_id=TASK_ID,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ mock_hook.assert_called_once_with(
+ api_version="v1",
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+ mock_hook.return_value.start_instance.assert_called_once_with(
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ project_id=None,
+ )
- def test_instance_start_should_throw_ex_when_missing_resource_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'resource_id' is missing"):
- ComputeEngineStartInstanceOperator(
- project_id=GCP_PROJECT_ID,
- zone=GCE_ZONE,
- resource_id="",
- task_id=TASK_ID,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_instance_start_should_throw_ex_when_missing_zone(self, mock_hook):
+ op = ComputeEngineStartInstanceOperator(
+ project_id=GCP_PROJECT_ID,
+ zone="",
+ resource_id=GCE_RESOURCE_ID,
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'zone' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
+
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_instance_start_should_throw_ex_when_missing_resource_id(self,
mock_hook):
+ op = ComputeEngineStartInstanceOperator(
+ project_id=GCP_PROJECT_ID,
+ zone=GCE_ZONE,
+ resource_id="",
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'resource_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
class TestGceInstanceStop:
@@ -655,14 +783,20 @@ class TestGceInstanceStop:
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.stop_instance.assert_called_once_with(
- zone=GCE_ZONE, resource_id=GCE_RESOURCE_ID,
project_id=GCP_PROJECT_ID
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ project_id=GCP_PROJECT_ID,
)
# Setting all the operator's input parameters as templated dag_ids
# (could be anything else) just to test if the templating works for all
fields
@pytest.mark.db_test
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
- def test_instance_stop_with_templates(self, _,
create_task_instance_of_operator, session):
+ def test_instance_stop_with_templates(
+ self,
+ _,
+ create_task_instance_of_operator,
+ ):
dag_id = "test_instance_stop_with_templates"
ti = create_task_instance_of_operator(
ComputeEngineStopInstanceOperator,
@@ -681,38 +815,78 @@ class TestGceInstanceStop:
assert dag_id == rendered.gcp_conn_id
assert dag_id == rendered.api_version
- def test_instance_stop_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineStopInstanceOperator(
- project_id="", zone=GCE_ZONE, resource_id=GCE_RESOURCE_ID,
task_id="id"
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_instance_stop_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineStopInstanceOperator(
+ project_id="",
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ task_id="id",
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_instance_stop_should_not_throw_ex_when_project_id_none(self,
mock_hook):
- op = ComputeEngineStopInstanceOperator(zone=GCE_ZONE,
resource_id=GCE_RESOURCE_ID, task_id="id")
+ op = ComputeEngineStopInstanceOperator(
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ task_id="id",
+ )
+
context = mock.MagicMock()
op.execute(context=context)
+
mock_hook.assert_called_once_with(
api_version="v1",
gcp_conn_id="google_cloud_default",
impersonation_chain=None,
)
mock_hook.return_value.stop_instance.assert_called_once_with(
- zone=GCE_ZONE, resource_id=GCE_RESOURCE_ID, project_id=None
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ project_id=None,
)
- def test_instance_stop_should_throw_ex_when_missing_zone(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'zone' is missing"):
- ComputeEngineStopInstanceOperator(
- project_id=GCP_PROJECT_ID, zone="",
resource_id=GCE_RESOURCE_ID, task_id="id"
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_instance_stop_should_throw_ex_when_missing_zone(self, mock_hook):
+ op = ComputeEngineStopInstanceOperator(
+ project_id=GCP_PROJECT_ID,
+ zone="",
+ resource_id=GCE_RESOURCE_ID,
+ task_id="id",
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'zone' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_instance_stop_should_throw_ex_when_missing_resource_id(self,
mock_hook):
- with pytest.raises(AirflowException, match=r"The required parameter
'resource_id' is missing"):
- ComputeEngineStopInstanceOperator(
- project_id=GCP_PROJECT_ID, zone=GCE_ZONE, resource_id="",
task_id="id"
- )
+ op = ComputeEngineStopInstanceOperator(
+ project_id=GCP_PROJECT_ID,
+ zone=GCE_ZONE,
+ resource_id="",
+ task_id="id",
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'resource_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
class TestGceInstanceSetMachineType:
@@ -736,14 +910,17 @@ class TestGceInstanceSetMachineType:
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.set_machine_type.assert_called_once_with(
- zone=GCE_ZONE, resource_id=GCE_RESOURCE_ID,
body=SET_MACHINE_TYPE_BODY, project_id=GCP_PROJECT_ID
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ body=SET_MACHINE_TYPE_BODY,
+ project_id=GCP_PROJECT_ID,
)
# Setting all the operator's input parameters as templated dag_ids
# (could be anything else) just to test if the templating works for all
fields
@pytest.mark.db_test
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
- def test_machine_type_set_with_templates(self, _,
create_task_instance_of_operator, session):
+ def test_machine_type_set_with_templates(self, _,
create_task_instance_of_operator):
dag_id = "test_set_machine_type_with_templates"
ti = create_task_instance_of_operator(
ComputeEngineSetMachineTypeOperator,
@@ -763,15 +940,23 @@ class TestGceInstanceSetMachineType:
assert dag_id == rendered.gcp_conn_id
assert dag_id == rendered.api_version
- def test_machine_type_set_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineSetMachineTypeOperator(
- project_id="",
- zone=GCE_ZONE,
- resource_id=GCE_RESOURCE_ID,
- body=SET_MACHINE_TYPE_BODY,
- task_id=TASK_ID,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_machine_type_set_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineSetMachineTypeOperator(
+ project_id="",
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ body=SET_MACHINE_TYPE_BODY,
+ task_id=TASK_ID,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_machine_type_set_should_not_throw_ex_when_project_id_none(self,
mock_hook):
@@ -783,39 +968,59 @@ class TestGceInstanceSetMachineType:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
context = mock.MagicMock()
op.execute(context=context)
+
mock_hook.assert_called_once_with(
api_version="v1",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.set_machine_type.assert_called_once_with(
- zone=GCE_ZONE, resource_id=GCE_RESOURCE_ID,
body=SET_MACHINE_TYPE_BODY, project_id=None
+ zone=GCE_ZONE,
+ resource_id=GCE_RESOURCE_ID,
+ body=SET_MACHINE_TYPE_BODY,
+ project_id=None,
)
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_machine_type_set_should_throw_ex_when_missing_zone(self,
mock_hook):
- with pytest.raises(AirflowException, match=r"The required parameter
'zone' is missing"):
- ComputeEngineSetMachineTypeOperator(
- project_id=GCP_PROJECT_ID,
- zone="",
- resource_id=GCE_RESOURCE_ID,
- body=SET_MACHINE_TYPE_BODY,
- task_id=TASK_ID,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ op = ComputeEngineSetMachineTypeOperator(
+ project_id=GCP_PROJECT_ID,
+ zone="",
+ resource_id=GCE_RESOURCE_ID,
+ body=SET_MACHINE_TYPE_BODY,
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
- def test_machine_type_set_should_throw_ex_when_missing_resource_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'resource_id' is missing"):
- ComputeEngineSetMachineTypeOperator(
- project_id=GCP_PROJECT_ID,
- zone=GCE_ZONE,
- resource_id="",
- body=SET_MACHINE_TYPE_BODY,
- task_id=TASK_ID,
- )
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'zone' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
+
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_machine_type_set_should_throw_ex_when_missing_resource_id(self,
mock_hook):
+ op = ComputeEngineSetMachineTypeOperator(
+ project_id=GCP_PROJECT_ID,
+ zone=GCE_ZONE,
+ resource_id="",
+ body=SET_MACHINE_TYPE_BODY,
+ task_id=TASK_ID,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'resource_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_machine_type_set_should_throw_ex_when_missing_machine_type(self,
mock_hook):
@@ -829,13 +1034,13 @@ class TestGceInstanceSetMachineType:
impersonation_chain=IMPERSONATION_CHAIN,
)
context = mock.MagicMock()
- with pytest.raises(AirflowException, match=r"The required body field
'machineType' is missing"):
+ with pytest.raises(
+ AirflowException,
+ match=r"The required body field 'machineType' is missing",
+ ):
op.execute(context=context)
- mock_hook.assert_called_once_with(
- api_version="v1",
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+
+ mock_hook.assert_not_called()
MOCK_OP_RESPONSE = (
"{'kind': 'compute#operation', 'id': '8529919847974922736', "
@@ -874,6 +1079,7 @@ class TestGceInstanceSetMachineType:
get_conn.return_value = {}
_execute_set_machine_type.return_value = {"name": "test-operation"}
_check_zone_operation_status.return_value =
ast.literal_eval(self.MOCK_OP_RESPONSE)
+
op = ComputeEngineSetMachineTypeOperator(
project_id=GCP_PROJECT_ID,
zone=GCE_ZONE,
@@ -881,14 +1087,27 @@ class TestGceInstanceSetMachineType:
body=SET_MACHINE_TYPE_BODY,
task_id=TASK_ID,
)
+
context = mock.MagicMock()
- with pytest.raises(AirflowException, match=r"400 BAD REQUEST:
{.+UNSUPPORTED_OPERATION"):
+
+ with pytest.raises(
+ AirflowException,
+ match=r"400 BAD REQUEST: {.+UNSUPPORTED_OPERATION",
+ ):
op.execute(context=context)
+
_check_zone_operation_status.assert_called_once_with(
- {}, "test-operation", GCP_PROJECT_ID, GCE_ZONE, mock.ANY
+ {},
+ "test-operation",
+ GCP_PROJECT_ID,
+ GCE_ZONE,
+ mock.ANY,
)
_execute_set_machine_type.assert_called_once_with(
- GCE_ZONE, GCE_RESOURCE_ID, SET_MACHINE_TYPE_BODY, GCP_PROJECT_ID
+ GCE_ZONE,
+ GCE_RESOURCE_ID,
+ SET_MACHINE_TYPE_BODY,
+ GCP_PROJECT_ID,
)
@@ -923,6 +1142,7 @@ class TestGceTemplateInsert:
NotFound("Error message"),
get_template_obj_mock,
]
+
op = ComputeEngineInsertInstanceTemplateOperator(
project_id=GCP_PROJECT_ID,
body=GCE_INSTANCE_TEMPLATE_BODY_API_CALL,
@@ -933,7 +1153,13 @@ class TestGceTemplateInsert:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
+ assert op.resource_id is None
+
op.execute(context=mock.MagicMock())
+
+ assert op.resource_id == GCE_INSTANCE_TEMPLATE_BODY_API_CALL["name"]
+
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
@@ -945,17 +1171,25 @@ class TestGceTemplateInsert:
request_id=None,
)
- def test_insert_template_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineInsertInstanceTemplateOperator(
- project_id="",
- body=GCE_INSTANCE_TEMPLATE_BODY_API_CALL,
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_insert_template_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineInsertInstanceTemplateOperator(
+ project_id="",
+ body=GCE_INSTANCE_TEMPLATE_BODY_API_CALL,
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_insert_template_should_not_throw_ex_when_project_id_none(self,
mock_hook):
@@ -965,6 +1199,7 @@ class TestGceTemplateInsert:
NotFound("Error message"),
get_template_obj_mock,
]
+
op = ComputeEngineInsertInstanceTemplateOperator(
body=GCE_INSTANCE_TEMPLATE_BODY_API_CALL,
task_id=TASK_ID,
@@ -974,7 +1209,13 @@ class TestGceTemplateInsert:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
+ assert op.resource_id is None
+
op.execute(context=mock.MagicMock())
+
+ assert op.resource_id == GCE_INSTANCE_TEMPLATE_BODY_API_CALL["name"]
+
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
@@ -1005,8 +1246,10 @@ class TestGceTemplateInsert:
NotFound("Error message"),
get_template_obj_mock,
]
+
body_with_templated_name =
deepcopy(GCE_INSTANCE_TEMPLATE_BODY_API_CALL)
body_with_templated_name["name"] = "{{ execution_date }}"
+
op = ComputeEngineInsertInstanceTemplateOperator(
project_id=GCP_PROJECT_ID,
body=body_with_templated_name,
@@ -1017,7 +1260,13 @@ class TestGceTemplateInsert:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
+ assert op.resource_id is None
+
op.execute(context=mock.MagicMock())
+
+ assert op.resource_id == body_with_templated_name["name"]
+
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
@@ -1055,18 +1304,26 @@ class TestGceTemplateDelete:
resource_id=GCE_RESOURCE_ID,
)
- def test_delete_template_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineDeleteInstanceTemplateOperator(
- project_id="",
- resource_id=GCE_RESOURCE_ID,
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_delete_template_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineDeleteInstanceTemplateOperator(
+ resource_id=GCE_RESOURCE_ID,
+ project_id="",
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_delete_template_should_not_throw_ex_when_project_id_none(self,
mock_hook):
@@ -1091,18 +1348,26 @@ class TestGceTemplateDelete:
request_id=None,
)
- def test_delete_template_should_throw_ex_when_missing_resource_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'resource_id' is missing"):
- ComputeEngineDeleteInstanceTemplateOperator(
- resource_id="",
- project_id=GCP_PROJECT_ID,
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_delete_template_should_throw_ex_when_missing_resource_id(self,
mock_hook):
+ op = ComputeEngineDeleteInstanceTemplateOperator(
+ resource_id="",
+ project_id=GCP_PROJECT_ID,
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'resource_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
GCE_INSTANCE_TEMPLATE_NAME = "instance-template-test"
@@ -1441,20 +1706,45 @@ class TestGceInstanceTemplateCopy:
request_id=None,
)
- def test_copy_template_with_missing_name_should_execute_successfully(self):
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_copy_template_should_throw_ex_when_missing_name(self, mock_hook):
+ op = ComputeEngineCopyInstanceTemplateOperator(
+ project_id=GCP_PROJECT_ID,
+ resource_id=GCE_INSTANCE_TEMPLATE_NAME,
+ request_id=GCE_INSTANCE_TEMPLATE_REQUEST_ID,
+ task_id=TASK_ID,
+ body_patch={"description": "New description"},
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
with pytest.raises(
AirflowException,
match=r"should contain at least name for the new operator in the
'name' field",
):
- ComputeEngineCopyInstanceTemplateOperator(
- project_id=GCP_PROJECT_ID,
- resource_id=GCE_INSTANCE_TEMPLATE_NAME,
- request_id=GCE_INSTANCE_TEMPLATE_REQUEST_ID,
- task_id=TASK_ID,
- body_patch={"description": "New description"},
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
+
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_copy_template_should_throw_ex_when_missing_resource_id(self,
mock_hook):
+ op = ComputeEngineCopyInstanceTemplateOperator(
+ project_id=GCP_PROJECT_ID,
+ resource_id="",
+ request_id=GCE_INSTANCE_TEMPLATE_REQUEST_ID,
+ task_id=TASK_ID,
+ body_patch={"name": "new-instance-template"},
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'resource_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
GCE_INSTANCE_GROUP_MANAGER_NAME = "instance-group-test"
@@ -1566,7 +1856,12 @@ class TestGceInstanceGroupManagerInsert:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
+ assert op.resource_id is None
+
op.execute(context=mock.MagicMock())
+
+ assert op.resource_id ==
GCE_INSTANCE_GROUP_MANAGER_BODY_API_CALL["name"]
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
@@ -1581,17 +1876,24 @@ class TestGceInstanceGroupManagerInsert:
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_insert_igm_should_throw_ex_when_missing_project_id(self,
mock_hook):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineInsertInstanceGroupManagerOperator(
- project_id="",
- body=GCE_INSTANCE_GROUP_MANAGER_BODY_API_CALL,
- zone=GCE_ZONE,
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ op = ComputeEngineInsertInstanceGroupManagerOperator(
+ project_id="",
+ body=GCE_INSTANCE_GROUP_MANAGER_BODY_API_CALL,
+ zone=GCE_ZONE,
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_insert_igm_should_not_throw_ex_when_project_id_none(self,
mock_hook):
@@ -1644,8 +1946,10 @@ class TestGceInstanceGroupManagerInsert:
NotFound("Error message"),
get_instance_group_manager_obj_mock,
]
+
body_with_templated_name =
deepcopy(GCE_INSTANCE_GROUP_MANAGER_BODY_API_CALL)
body_with_templated_name["name"] = "{{ execution_date }}"
+
op = ComputeEngineInsertInstanceGroupManagerOperator(
project_id=GCP_PROJECT_ID,
body=body_with_templated_name,
@@ -1657,7 +1961,13 @@ class TestGceInstanceGroupManagerInsert:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
+ assert op.resource_id is None
+
op.execute(context=mock.MagicMock())
+
+ assert op.resource_id == body_with_templated_name["name"]
+
mock_hook.assert_called_once_with(
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
@@ -1698,19 +2008,27 @@ class TestGceInstanceGroupManagerDelete:
project_id=GCP_PROJECT_ID,
)
- def test_delete_igm_should_throw_ex_when_missing_project_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'project_id' is missing"):
- ComputeEngineDeleteInstanceGroupManagerOperator(
- project_id="",
- resource_id=GCE_INSTANCE_GROUP_MANAGER_NAME,
- zone=GCE_ZONE,
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_delete_igm_should_throw_ex_when_missing_project_id(self,
mock_hook):
+ op = ComputeEngineDeleteInstanceGroupManagerOperator(
+ project_id="",
+ resource_id=GCE_INSTANCE_GROUP_MANAGER_NAME,
+ zone=GCE_ZONE,
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'project_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
def test_delete_igm_should_not_throw_ex_when_project_id_none(self,
mock_hook):
@@ -1737,18 +2055,27 @@ class TestGceInstanceGroupManagerDelete:
project_id=None,
)
- def test_delete_igm_should_throw_ex_when_missing_resource_id(self):
- with pytest.raises(AirflowException, match=r"The required parameter
'resource_id' is missing"):
- ComputeEngineDeleteInstanceGroupManagerOperator(
- resource_id="",
- zone=GCE_ZONE,
- task_id=TASK_ID,
- retry=RETRY,
- timeout=TIMEOUT,
- metadata=METADATA,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
+ @mock.patch(COMPUTE_ENGINE_HOOK_PATH)
+ def test_delete_igm_should_throw_ex_when_missing_resource_id(self,
mock_hook):
+ op = ComputeEngineDeleteInstanceGroupManagerOperator(
+ project_id=GCP_PROJECT_ID,
+ zone=GCE_ZONE,
+ resource_id="",
+ task_id=TASK_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"The required parameter 'resource_id' is missing",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
class TestGceInstanceGroupManagerUpdate:
@@ -1773,7 +2100,9 @@ class TestGceInstanceGroupManagerUpdate:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
result = op.execute(context=mock.MagicMock())
+
mock_hook.assert_called_once_with(
api_version="beta",
gcp_conn_id=GCP_CONN_ID,
@@ -1797,6 +2126,7 @@ class TestGceInstanceGroupManagerUpdate:
get_instance_group_manager_obj_mock,
]
igm.to_dict.return_value = GCE_INSTANCE_GROUP_MANAGER_GET
+
op = ComputeEngineInstanceGroupUpdateManagerTemplateOperator(
zone=GCE_ZONE,
resource_id=GCE_INSTANCE_GROUP_MANAGER_NAME,
@@ -1807,7 +2137,9 @@ class TestGceInstanceGroupManagerUpdate:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
result = op.execute(context=mock.MagicMock())
+
mock_hook.assert_called_once_with(
api_version="beta",
gcp_conn_id=GCP_CONN_ID,
@@ -1829,12 +2161,14 @@ class TestGceInstanceGroupManagerUpdate:
):
instance_group_manager_no_template =
deepcopy(GCE_INSTANCE_GROUP_MANAGER_GET)
del instance_group_manager_no_template["instanceTemplate"]
+
get_instance_group_manager_obj_mock = mock.MagicMock()
get_instance_group_manager_obj_mock.__class__ = InstanceGroupManager
mock_hook.return_value.get_instance_group_manager.side_effect = [
get_instance_group_manager_obj_mock,
]
igm.to_dict.return_value = instance_group_manager_no_template
+
op = ComputeEngineInstanceGroupUpdateManagerTemplateOperator(
project_id=GCP_PROJECT_ID,
zone=GCE_ZONE,
@@ -1846,14 +2180,18 @@ class TestGceInstanceGroupManagerUpdate:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
result = op.execute(context=mock.MagicMock())
+
mock_hook.assert_called_once_with(
api_version="beta",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
expected_patch_no_instance_template =
deepcopy(GCE_INSTANCE_GROUP_MANAGER_EXPECTED_PATCH)
del expected_patch_no_instance_template["instanceTemplate"]
+
mock_hook.return_value.patch_instance_group_manager.assert_called_once_with(
project_id=GCP_PROJECT_ID,
zone=GCE_ZONE,
@@ -1868,12 +2206,14 @@ class TestGceInstanceGroupManagerUpdate:
def
test_update_instance_group_no_versions_field_should_execute_successfully(self,
mock_hook, igm):
instance_group_manager_no_versions =
deepcopy(GCE_INSTANCE_GROUP_MANAGER_GET)
del instance_group_manager_no_versions["versions"]
+
get_instance_group_manager_obj_mock = mock.MagicMock()
get_instance_group_manager_obj_mock.__class__ = InstanceGroupManager
mock_hook.return_value.get_instance_group_manager.side_effect = [
get_instance_group_manager_obj_mock,
]
igm.to_dict.return_value = instance_group_manager_no_versions
+
op = ComputeEngineInstanceGroupUpdateManagerTemplateOperator(
project_id=GCP_PROJECT_ID,
zone=GCE_ZONE,
@@ -1885,14 +2225,18 @@ class TestGceInstanceGroupManagerUpdate:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
result = op.execute(context=mock.MagicMock())
+
mock_hook.assert_called_once_with(
api_version="beta",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+
expected_patch_no_versions =
deepcopy(GCE_INSTANCE_GROUP_MANAGER_EXPECTED_PATCH)
del expected_patch_no_versions["versions"]
+
mock_hook.return_value.patch_instance_group_manager.assert_called_once_with(
project_id=GCP_PROJECT_ID,
zone=GCE_ZONE,
@@ -1939,17 +2283,24 @@ class TestGceInstanceGroupManagerUpdate:
assert result
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
- def test_update_instance_group_try_to_use_api_v1_should_throw_ex(self, _):
- with pytest.raises(AirflowException, match=r"Use beta api version or
above"):
- ComputeEngineInstanceGroupUpdateManagerTemplateOperator(
- project_id=GCP_PROJECT_ID,
- zone=GCE_ZONE,
- resource_id=GCE_INSTANCE_GROUP_MANAGER_NAME,
- task_id=TASK_ID,
- api_version="v1",
- source_template=GCE_INSTANCE_TEMPLATE_SOURCE_URL,
- destination_template=GCE_INSTANCE_TEMPLATE_DESTINATION_URL,
- )
+ def test_update_instance_group_try_to_use_api_v1_should_throw_ex(self,
mock_hook):
+ op = ComputeEngineInstanceGroupUpdateManagerTemplateOperator(
+ project_id=GCP_PROJECT_ID,
+ zone=GCE_ZONE,
+ resource_id=GCE_INSTANCE_GROUP_MANAGER_NAME,
+ task_id=TASK_ID,
+ api_version="v1",
+ source_template=GCE_INSTANCE_TEMPLATE_SOURCE_URL,
+ destination_template=GCE_INSTANCE_TEMPLATE_DESTINATION_URL,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=r"Use beta api version or above",
+ ):
+ op.execute(context=mock.MagicMock())
+
+ mock_hook.assert_not_called()
@mock.patch(IGM_PATH)
@mock.patch(COMPUTE_ENGINE_HOOK_PATH)
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index dbe391c48fc..378be2590c7 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -23,16 +23,6 @@
providers/google/src/airflow/providers/google/cloud/operators/bigquery.py::BigQu
providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py::CloudBatchSubmitJobOperator
providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py::CloudBuildCreateBuildOperator
providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py::CloudDataTransferServiceCreateJobOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineCopyInstanceTemplateOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineDeleteInstanceGroupManagerOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineDeleteInstanceOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineDeleteInstanceTemplateOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceFromTemplateOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceGroupManagerOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceTemplateOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInstanceGroupUpdateManagerTemplateOperator
-providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineSetMachineTypeOperator
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocCreateClusterOperator
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocSubmitJobOperator
providers/google/src/airflow/providers/google/cloud/operators/functions.py::CloudFunctionDeployFunctionOperator
