def __init__(self, step_kind, step_name, additional_properties=None):
self.step_kind = step_kind
self.step_name = step_name
self.proto = dataflow.Step(kind=step_kind, name=step_name)
self.proto.properties = {}
self._additional_properties = []
if additional_properties is not None:
for (n, v, t) in additional_properties:
self.add_property(n, v, t)
def test_template_file_generation_with_upload_graph(self):
pipeline_options = PipelineOptions([
'--project',
'test_project',
'--job_name',
'test_job_name',
'--temp_location',
'gs://test-location/temp',
'--experiments',
'upload_graph',
'--template_location',
'gs://test-location/template'
])
job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL)
job.proto.steps.append(dataflow.Step(name='test_step_name'))
pipeline_options.view_as(GoogleCloudOptions).no_auth = True
client = apiclient.DataflowApplicationClient(pipeline_options)
with mock.patch.object(client, 'stage_file', side_effect=None):
with mock.patch.object(client, 'create_job_description',
side_effect=None):
with mock.patch.object(client,
'submit_job_description',
side_effect=None):
client.create_job(job)
client.stage_file.assert_has_calls([
mock.call(mock.ANY, 'dataflow_graph.json', mock.ANY),
mock.call(mock.ANY, 'template', mock.ANY)
])
client.create_job_description.assert_called_once()
# template is generated, but job should not be submitted to the
# service.
client.submit_job_description.assert_not_called()
template_filename = client.stage_file.call_args_list[-1][0][1]
self.assertTrue('template' in template_filename)
template_content = client.stage_file.call_args_list[-1][0][2].read(
).decode('utf-8')
template_obj = json.loads(template_content)
self.assertFalse(template_obj.get('steps'))
self.assertTrue(template_obj['stepsLocation'])
