def test_save_as(self):
params = SaveResultsRequestParams()
params.file_path = 'somepath'
params.row_start_index = 0
params.row_end_index = 1
mock_writer = MockWriter(10)
mock_file_factory = mock.MagicMock()
mock_file_factory.get_writer = mock.Mock(return_value=mock_writer)
on_success = mock.MagicMock()
self._result_set._has_been_read = True
self._result_set.rows.append(self._first_row)
self._result_set.get_row = mock.Mock(return_value=self._first_row)
self._result_set.save_as(params, mock_file_factory, on_success, None)
mock_file_factory.get_writer.assert_called_once_with(params.file_path)
mock_writer.write_row.assert_called_once_with(
self._first_row, self._result_set.columns_info)
mock_writer.complete_write.assert_called_once()
on_success.assert_called_once()
def execute_with_batch_index(index: int):
params = SaveResultsRequestParams()
params.batch_index = index
with self.assertRaises(IndexError) as context_manager:
self.query.save_as(params, None, None, None)
self.assertEquals(
'Batch index cannot be less than 0 or greater than the number of batches',
context_manager.exception.args[0])
def test_save_as_with_invalid_batch_index(self):
batch = self.create_and_execute_batch(Batch)
params = SaveResultsRequestParams()
params.result_set_index = 1
with self.assertRaises(IndexError) as context_manager:
batch.save_as(params, None, None, None)
self.assertEquals('Result set index should be always 0',
context_manager.exception.args[0])
def test_save_as_result_set_when_thread_alive(self):
params = SaveResultsRequestParams()
params.file_path = 'some_file_path'
mock_thread = mock.MagicMock()
mock_thread.isAlive = True
self._result_set._save_as_threads[params.file_path] = mock_thread
with self.assertRaises(RuntimeError) as context_manager:
self._result_set.save_as(params, None, None, None)
self.assertEquals('A save request to the same path is in progress',
context_manager.exception.args[0])
def test_save_as(self):
params = SaveResultsRequestParams()
params.batch_index = 0
file_factory = mock.MagicMock()
on_success = mock.MagicMock()
on_error = mock.MagicMock()
batch_save_as_mock = mock.MagicMock()
self.query.batches[0].save_as = batch_save_as_mock
self.query.save_as(params, file_factory, on_success, on_error)
batch_save_as_mock.assert_called_once_with(params, file_factory,
on_success, on_error)
def test_save_as_result_set_when_not_read(self):
params = SaveResultsRequestParams()
with self.assertRaises(RuntimeError) as context_manager:
self._result_set.save_as(params, None, None, None)
self.assertEquals(
'Result cannot be saved until query execution has completed',
context_manager.exception.args[0])
def test_save_as(self):
batch = self.create_and_execute_batch(Batch)
params = SaveResultsRequestParams()
params.result_set_index = 0
file_factory = mock.MagicMock()
on_success = mock.MagicMock()
on_error = mock.MagicMock()
result_set_save_as_mock = mock.MagicMock()
batch._result_set.save_as = result_set_save_as_mock
batch.save_as(params, file_factory, on_success, on_error)
result_set_save_as_mock.assert_called_once_with(
params, file_factory, on_success, on_error)
def setUp(self):
self.request = SaveResultsRequestParams()
self.request.file_path = 'TestPath'
self.factory = SaveAsJsonFileStreamFactory(self.request)