class TestDataflowTemplateOperator(unittest.TestCase): def setUp(self): self.dataflow = DataflowTemplatedJobStartOperator( task_id=TASK_ID, template=TEMPLATE, job_name=JOB_NAME, parameters=PARAMETERS, dataflow_default_options=DEFAULT_OPTIONS_TEMPLATE, poll_sleep=POLL_SLEEP, location=TEST_LOCATION) def test_init(self): """Test DataflowTemplateOperator instance is properly initialized.""" self.assertEqual(self.dataflow.task_id, TASK_ID) self.assertEqual(self.dataflow.job_name, JOB_NAME) self.assertEqual(self.dataflow.template, TEMPLATE) self.assertEqual(self.dataflow.parameters, PARAMETERS) self.assertEqual(self.dataflow.poll_sleep, POLL_SLEEP) self.assertEqual(self.dataflow.dataflow_default_options, DEFAULT_OPTIONS_TEMPLATE) @mock.patch( 'airflow.providers.google.cloud.operators.dataflow.DataflowHook') def test_exec(self, dataflow_mock): """Test DataflowHook is created and the right args are passed to start_template_workflow. """ start_template_hook = dataflow_mock.return_value.start_template_dataflow self.dataflow.execute(None) self.assertTrue(dataflow_mock.called) expected_options = { 'project': 'test', 'stagingLocation': 'gs://test/staging', 'tempLocation': 'gs://test/temp', 'zone': 'us-central1-f' } start_template_hook.assert_called_once_with( job_name=JOB_NAME, variables=expected_options, parameters=PARAMETERS, dataflow_template=TEMPLATE, on_new_job_id_callback=mock.ANY, project_id=None, location=TEST_LOCATION)
class TestDataflowTemplateOperator(unittest.TestCase): def setUp(self): self.dataflow = DataflowTemplatedJobStartOperator( task_id=TASK_ID, template=TEMPLATE, job_name=JOB_NAME, parameters=PARAMETERS, options=DEFAULT_OPTIONS_TEMPLATE, dataflow_default_options={"EXTRA_OPTION": "TEST_A"}, poll_sleep=POLL_SLEEP, location=TEST_LOCATION, environment={"maxWorkers": 2}, ) @mock.patch( 'airflow.providers.google.cloud.operators.dataflow.DataflowHook') def test_exec(self, dataflow_mock): """Test DataflowHook is created and the right args are passed to start_template_workflow. """ start_template_hook = dataflow_mock.return_value.start_template_dataflow self.dataflow.execute(None) self.assertTrue(dataflow_mock.called) expected_options = { 'project': 'test', 'stagingLocation': 'gs://test/staging', 'tempLocation': 'gs://test/temp', 'zone': 'us-central1-f', 'EXTRA_OPTION': "TEST_A", } start_template_hook.assert_called_once_with( job_name=JOB_NAME, variables=expected_options, parameters=PARAMETERS, dataflow_template=TEMPLATE, on_new_job_id_callback=mock.ANY, project_id=None, location=TEST_LOCATION, environment={'maxWorkers': 2}, )