def test_default_parameters_extract_operator(self):
extract_operator = ExtractOperator(task_id='foo_task_id',
path='foo_path',
table='foo_table')
self.assertEquals(extract_operator._format, 'csv')
self.assertDictEqual(extract_operator._csv_options, {})
self.assertEqual(extract_operator.template_fields, ('_table', '_path'))
def test_extract_operator(self, mock_table_extract, mock_context_default):
mock_context_default.return_value = TestCases._create_context()
extract_operator = ExtractOperator(table=TestCases.test_project_id +
'.test_table',
path='test_path',
format=None,
task_id='test_extract_operator')
# Happy path
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = False
mock_table_extract.return_value.errors = None
self.assertDictEqual(extract_operator.execute(context=None),
{'result': 'test-results'})
mock_table_extract.assert_called_with('test_path',
format='NEWLINE_DELIMITED_JSON',
csv_delimiter=None,
csv_header=None,
compress=None)
# Extract failed
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = True
mock_table_extract.return_value.errors = None
with self.assertRaisesRegexp(Exception, "Extract failed:"):
extract_operator.execute(context=None)
# Extract completed with errors
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = False
mock_table_extract.return_value.errors = 'foo_error'
with self.assertRaisesRegexp(
Exception, 'Extract completed with errors: foo_error'):
extract_operator.execute(context=None)
def test_extract_operator(self, mock_table_extract, mock_context_default):
mock_context_default.return_value = TestCases._create_context()
extract_operator = ExtractOperator(table=TestCases.test_project_id +
'.test_table',
path='test_path',
format=None,
task_id='test_extract_operator')
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = False
mock_table_extract.return_value.errors = None
self.assertEqual(extract_operator.execute(context=None),
'test-results')
mock_table_extract.assert_called_with('test_path',
format='NEWLINE_DELIMITED_JSON',
csv_delimiter=None,
csv_header=None,
compress=None)
def test_extract_operator(self, mock_table_extract, mock_context_default):
mock_context_default.return_value = TestCases._create_context()
extract_operator = ExtractOperator(table=TestCases.test_project_id + '.test_table',
path='test_path', format=None,
task_id='test_extract_operator',
csv_options={'delimiter': '$'})
# Happy path
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = False
mock_table_extract.return_value.errors = None
self.assertDictEqual(extract_operator.execute(context=None), {'result': 'test-results'})
mock_table_extract.assert_called_with('test_path', format='NEWLINE_DELIMITED_JSON',
csv_delimiter='$')
# Extract failed
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = True
mock_table_extract.return_value.errors = None
with self.assertRaisesRegexp(Exception, "Extract failed:"):
extract_operator.execute(context=None)
# Extract completed with errors
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = False
mock_table_extract.return_value.errors = 'foo_error'
with self.assertRaisesRegexp(Exception, 'Extract completed with errors: foo_error'):
extract_operator.execute(context=None)
def test_extract_operator_with_temporary_table(self, mock_task_instance,
mock_table_extract,
mock_context_default):
mock_context_default.return_value = TestCases._create_context()
mock_task_instance.xcom_pull.return_value = {
'table': TestCases.test_project_id + '.test_table'
}
extract_operator = ExtractOperator(path='test_path',
format=None,
task_id='test_extract_operator')
mock_table_extract.return_value.result = lambda: 'test-results'
mock_table_extract.return_value.failed = False
mock_table_extract.return_value.errors = None
self.assertDictEqual(
extract_operator.execute(
context={'task_instance': mock_task_instance}),
{'result': 'test-results'})
mock_table_extract.assert_called_with('test_path',
format='NEWLINE_DELIMITED_JSON',
csv_delimiter=None,
csv_header=None,
compress=None)