def test_execute_operator(self, mock_query_job, mock_query_output_table, mock_query_execute,
mock_context_default):
mock_context_default.return_value = self._create_context()
execute_operator = ExecuteOperator(task_id='test_execute_operator', sql='test_sql')
mock_query_execute.return_value = mock_query_job
mock_query_job.result.return_value = google.datalab.bigquery.Table(TestCases.test_table_name)
self.assertDictEqual(execute_operator.execute(context=None),
{'table': TestCases.test_table_name})
mock_query_output_table.assert_called_with(name=None, use_cache=False,
allow_large_results=False)
def test_execute_operator(self, mock_get_notebook_item, mock_query_execute,
mock_context_default):
mock_context_default.return_value = self._create_context()
mock_get_notebook_item.return_value = google.datalab.bigquery.Query(
'test_sql')
# This statement is required even though it seems like it's not. Go figure.
execute_operator = ExecuteOperator(task_id='test_execute_operator',
sql='test_sql',
parameters=None,
table='test_table',
mode=None)
execute_operator.execute(context=None)
# TODO(rajivpb): Mock output_options and query_params for a more complete test.
mock_query_execute.assert_called_once()
def test_execute_operator(self, mock_query_job, mock_query_output_table,
mock_query_execute, mock_context_default):
mock_context_default.return_value = self._create_context()
execute_operator = ExecuteOperator(task_id='test_execute_operator',
sql='test_sql')
mock_query_execute.return_value = mock_query_job
query_results_table_name = 'foo_table'
mock_query_job.result.return_value.name = query_results_table_name
self.assertDictEqual(execute_operator.execute(context=None),
{'table': query_results_table_name})
mock_query_output_table.assert_called_with(name=None,
mode=None,
use_cache=False,
allow_large_results=False)
def test_execute_operator(self, mock_query_job, mock_query_output_table,
mock_query_execute, mock_context_default):
mock_context_default.return_value = self._create_context()
execute_operator = ExecuteOperator(task_id='test_execute_operator',
sql='test_sql')
mock_query_execute.return_value = mock_query_job
mock_query_job.result.return_value = google.datalab.bigquery.Table(
TestCases.test_table_name)
self.assertDictEqual(execute_operator.execute(context=None),
{'table': TestCases.test_table_name})
mock_query_output_table.assert_called_with(name=None,
mode=None,
use_cache=False,
allow_large_results=False)
def test_defaults_execute_operator(self):
execute_operator = ExecuteOperator(task_id='foo_task_id',
sql='foo_sql')
self.assertIsNone(execute_operator._parameters)
self.assertIsNone(execute_operator._table)
self.assertIsNone(execute_operator._mode)
self.assertEqual(execute_operator.template_fields, ('_sql', '_table'))
def test_execute_operator_defaults(self):
execute_operator = ExecuteOperator(task_id='foo_task_id',
sql='foo_sql')
self.assertIsNone(execute_operator.parameters)
self.assertIsNone(execute_operator.table)
self.assertIsNone(execute_operator.mode)
self.assertEqual(execute_operator.template_fields,
('table', 'parameters', 'path'))
def test_execute_operator_with_data_source(self, mock_query_job,
mock_query_output_table,
mock_query_class,
mock_external_data_source,
mock_context_default):
mock_context_default.return_value = self._create_context()
csv_options = {
'delimiter': 'f',
'skip': 9,
'strict': True,
'quote': 'l'
}
execute_operator = ExecuteOperator(task_id='test_execute_operator',
sql='test_sql',
data_source='foo_data_source',
path='foo_path',
max_bad_records=20,
schema=TestCases.test_schema,
csv_options=csv_options,
format='csv')
mock_query_instance = mock_query_class.return_value
mock_query_instance.execute.return_value = mock_query_job
mock_query_job.result.return_value = google.datalab.bigquery.Table(
TestCases.test_table_name)
self.assertDictEqual(execute_operator.execute(context=None),
{'table': TestCases.test_table_name})
mock_query_output_table.assert_called_with(name=None,
mode=None,
use_cache=False,
allow_large_results=False)
mock_query_class.assert_called_with(
sql='test_sql',
data_sources={
'foo_data_source': mock_external_data_source.return_value
})
mock_external_data_source.assert_called_with(
source='foo_path',
max_bad_records=20,
csv_options=mock.ANY,
source_format='csv',
schema=google.datalab.bigquery.Schema(TestCases.test_schema))
def test_execute_operator_with_data_source(self, mock_query_job, mock_query_output_table,
mock_query_class, mock_external_data_source,
mock_context_default):
mock_context_default.return_value = self._create_context()
csv_options = {'delimiter': 'f', 'skip': 9, 'strict': True, 'quote': 'l'}
execute_operator = ExecuteOperator(task_id='test_execute_operator', sql='test_sql',
data_source='foo_data_source', path='foo_path',
max_bad_records=20, schema=TestCases.test_schema,
csv_options=csv_options, format='csv')
mock_query_instance = mock_query_class.return_value
mock_query_instance.execute.return_value = mock_query_job
mock_query_job.result.return_value = google.datalab.bigquery.Table(TestCases.test_table_name)
self.assertDictEqual(execute_operator.execute(context=None),
{'table': TestCases.test_table_name})
mock_query_output_table.assert_called_with(name=None, use_cache=False,
allow_large_results=False)
mock_query_class.assert_called_with(
sql='test_sql', data_sources={'foo_data_source': mock_external_data_source.return_value})
mock_external_data_source.assert_called_with(
source='foo_path', max_bad_records=20, csv_options=mock.ANY, source_format='csv',
schema=google.datalab.bigquery.Schema(TestCases.test_schema))
execute_operator = ExecuteOperator(task_id='test_execute_operator', sql='test_sql',
data_source='foo_data_source', path='foo_path',
schema=TestCases.test_schema)
mock_query_instance = mock_query_class.return_value
mock_query_instance.execute.return_value = mock_query_job
execute_operator.execute(None)
mock_query_output_table.assert_called_with(name=None, use_cache=False,
allow_large_results=False)
mock_query_class.assert_called_with(sql='test_sql',
data_sources={'foo_data_source':
mock_external_data_source.return_value})
mock_external_data_source.assert_called_with(source='foo_path',
schema=google.datalab.bigquery.Schema(
TestCases.test_schema))