def execute(self, context):
hook = DataFlowHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
poll_sleep=self.poll_sleep)
hook.start_template_dataflow(self.job_name, self.dataflow_default_options,
self.parameters, self.template)
class TestDataFlowTemplateHook(unittest.TestCase):
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.dataflow_hook = DataFlowHook(gcp_conn_id='test')
@mock.patch(DATAFLOW_STRING.format('DataFlowHook._start_template_dataflow')
)
def test_start_template_dataflow(self, internal_dataflow_mock):
self.dataflow_hook.start_template_dataflow(
job_name=JOB_NAME,
variables=DATAFLOW_OPTIONS_TEMPLATE,
parameters=PARAMETERS,
dataflow_template=TEMPLATE)
options_with_region = {'region': 'us-central1'}
options_with_region.update(DATAFLOW_OPTIONS_TEMPLATE)
options_with_region_without_project = copy.deepcopy(
options_with_region)
del options_with_region_without_project['project']
internal_dataflow_mock.assert_called_once_with(
mock.ANY, options_with_region_without_project, PARAMETERS,
TEMPLATE, DATAFLOW_OPTIONS_JAVA['project'])
@mock.patch(DATAFLOW_STRING.format('_DataflowJobsController'))
@mock.patch(DATAFLOW_STRING.format('DataFlowHook.get_conn'))
def test_start_template_dataflow_with_runtime_env(self, mock_conn,
mock_dataflowjob):
dataflow_options_template = copy.deepcopy(DATAFLOW_OPTIONS_TEMPLATE)
options_with_runtime_env = copy.deepcopy(RUNTIME_ENV)
options_with_runtime_env.update(dataflow_options_template)
dataflowjob_instance = mock_dataflowjob.return_value
dataflowjob_instance.wait_for_done.return_value = None
method = (mock_conn.return_value.projects.return_value.locations.
return_value.templates.return_value.launch)
self.dataflow_hook.start_template_dataflow(
job_name=JOB_NAME,
variables=options_with_runtime_env,
parameters=PARAMETERS,
dataflow_template=TEMPLATE)
body = {
"jobName": mock.ANY,
"parameters": PARAMETERS,
"environment": RUNTIME_ENV
}
method.assert_called_once_with(
projectId=options_with_runtime_env['project'],
location='us-central1',
gcsPath=TEMPLATE,
body=body,
)