def test_create_inspect_template(self, mock_hook):
mock_hook.return_value.create_inspect_template.return_value = {}
operator = CloudDLPCreateInspectTemplateOperator(
organization_id=ORGANIZATION_ID, task_id="id")
operator.execute(context=None)
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
mock_hook.return_value.create_inspect_template.assert_called_once_with(
organization_id=ORGANIZATION_ID,
project_id=None,
inspect_template=None,
template_id=None,
retry=None,
timeout=None,
metadata=None,
)
}]
}],
})
INSPECT_CONFIG = InspectConfig(info_types=[{
"name": "PHONE_NUMBER"
}, {
"name": "US_TOLLFREE_PHONE_NUMBER"
}])
INSPECT_TEMPLATE = InspectTemplate(inspect_config=INSPECT_CONFIG)
with DAG("example_gcp_dlp", default_args=default_args,
schedule_interval=None) as dag:
create_template = CloudDLPCreateInspectTemplateOperator(
project_id=GCP_PROJECT,
inspect_template=INSPECT_TEMPLATE,
template_id=TEMPLATE_ID,
task_id="create_template",
do_xcom_push=True,
dag=dag,
)
inspect_content = CloudDLPInspectContentOperator(
task_id="inpsect_content",
project_id=GCP_PROJECT,
item=ITEM,
inspect_template_name=
"{{ task_instance.xcom_pull('create_template', key='return_value')['name'] }}",
dag=dag,
)
delete_template = CloudDLPDeleteInspectTemplateOperator(
task_id="delete_template",