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))