def test_filter_field_not_required_with_one_worker(self):
_config = {
'sql_bind': 'sql_bind',
'query': 'query',
'no_of_workers': 1
}
e = SQLExtractor(_config)
e.validate_config()
def test_no_filter_field_with_more_worker(self):
_config = {
'sql_bind': 'sql_bind',
'query': 'query',
'no_of_workers': 2
}
e = SQLExtractor(_config)
with pytest.raises(MissingConfigError,
match=r'.* filter_field .* SQL'):
e.validate_config()
def test_sync_extract_call_for_single_worker(self, export_to_parquet,
mock_uuid):
_valid_config['no_of_workers'] = 1
e = SQLExtractor(_valid_config)
export_to_parquet.return_value = "file1.parq"
self.assertEqual(e.extract_to_parquet(), ["file1.parq"])
export_to_parquet.assert_called_with(worker_id=0,
filter_field='filter_field',
output_folder='./F43C2651',
query='query %s',
sql_bind='sql_bind',
start_pos=None,
end_pos=None,
table_schema=None)
def test_async_extract_call(self, async_worker, get_results_cursor,
mock_uuid):
_valid_config['no_of_workers'] = 2
e = SQLExtractor(_valid_config)
mock_results = MagicMock()
get_results_cursor.return_value = mock_results
mock_results.fetchone.return_value = [3, 6]
mock_worker = MagicMock()
async_worker.return_value = mock_worker
mock_worker.send_data_to_worker = MagicMock()
mock_job_results = MagicMock(return_value=["file1.parq"])
mock_worker.get_job_results = mock_job_results
self.assertEqual(e.extract_to_parquet(), ["file1.parq"])
async_worker.assert_called_with(2)
mock_worker.send_data_to_worker.assert_called()
_call_args = mock_worker.send_data_to_worker.call_args_list
_, kwargs1 = _call_args[0]
self.assertEqual(kwargs1['worker_id'], 0)
self.assertEqual(kwargs1['worker_callback'].__name__,
'export_to_parquet')
self.assertEqual(kwargs1['sql_bind'], 'sql_bind')
self.assertEqual(kwargs1['query'], 'query %s')
self.assertEqual(kwargs1['filter_field'], 'filter_field')
self.assertEqual(kwargs1['output_folder'], './F43C2651')
self.assertEqual(kwargs1['start_pos'], 3)
self.assertEqual(kwargs1['end_pos'], 4)
_, kwargs2 = _call_args[1]
self.assertEqual(kwargs2['worker_id'], 1)
self.assertEqual(kwargs2['worker_callback'].__name__,
'export_to_parquet')
self.assertEqual(kwargs2['sql_bind'], 'sql_bind')
self.assertEqual(kwargs2['query'], 'query %s')
self.assertEqual(kwargs2['filter_field'], 'filter_field')
self.assertEqual(kwargs2['output_folder'], './F43C2651')
self.assertEqual(kwargs2['start_pos'], 5)
self.assertEqual(kwargs2['end_pos'], 6)
mock_job_results.assert_called_with()
def test_default_confgs(self, mock_uuid):
e = SQLExtractor(_valid_config)
self.assertEqual(e._no_of_workers, 2)
self.assertEqual(e._output_folder, './F43C2651')
def test_no_index(self):
_config = {'sql_bind': 'sql_bind'}
e = SQLExtractor(_config)
with pytest.raises(MissingConfigError, match=r'.* query .* SQL'):
e.validate_config()
def test_no_sql_bind(self):
_config = {}
e = SQLExtractor(_config)
with pytest.raises(MissingConfigError, match=r'.* sql_bind .* SQL'):
e.validate_config()
def test_with_minimum_required_configs(self):
e = SQLExtractor(_valid_config)
# No error raised
e.validate_config()