def test_assert_valid_hook_call(self, mock_hook):
task = CloudMemorystoreCreateInstanceOperator(
task_id=TEST_TASK_ID,
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
instance=TEST_INSTANCE,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
task.execute(mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.create_instance.assert_called_once_with(
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
instance=TEST_INSTANCE,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
FIRST_INSTANCE = {"tier": Instance.Tier.BASIC, "memory_size_gb": 1}
# [END howto_operator_instance]
SECOND_INSTANCE = {"tier": Instance.Tier.STANDARD_HA, "memory_size_gb": 3}
with models.DAG(
"gcp_cloud_memorystore",
schedule_interval=None, # Override to match your needs
start_date=dates.days_ago(1),
tags=['example'],
) as dag:
# [START howto_operator_create_instance]
create_instance = CloudMemorystoreCreateInstanceOperator(
task_id="create-instance",
location="europe-north1",
instance_id=INSTANCE_NAME,
instance=FIRST_INSTANCE,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_create_instance]
# [START howto_operator_create_instance_result]
create_instance_result = BashOperator(
task_id="create-instance-result",
bash_command=
"echo \"{{ task_instance.xcom_pull('create-instance') }}\"",
)
# [END howto_operator_create_instance_result]
create_instance_2 = CloudMemorystoreCreateInstanceOperator(
task_id="create-instance-2",