def test_get_subset(self):
full_query = 'Select * from t1;'
query = Query('test_uri', full_query,
QueryExecutionSettings(ExecutionPlanOptions(), None),
QueryEvents())
expected_subset = []
mock_batch = mock.MagicMock()
mock_batch.get_subset = mock.Mock(return_value=expected_subset)
query._batches = [mock_batch]
subset = query.get_subset(0, 0, 10)
self.assertEqual(expected_subset, subset)
mock_batch.get_subset.assert_called_once_with(0, 10)
def test_initialize_calls_success(self):
query = Query('owner', '', QueryExecutionSettings(None, None),
QueryEvents())
query._execution_state = ExecutionState.EXECUTED
rows = [("Result1", 53), (
"Result2",
None,
)]
result_set = self.get_result_set(rows)
batch = Batch('', 1, None)
batch._result_set = result_set
query._batches = [batch]
self._query_executer = mock.MagicMock(
return_value=DataEditSessionExecutionState(query))
self._data_editor_session.initialize(self._initialize_edit_request,
self._connection,
self._query_executer,
self._on_success,
self._on_failure)
self._query_executer.assert_called_once()