async def test_scroll_error(
aes,
index_name,
):
aes = MagicMock(spec=AsyncElasticsearch)
_search_response = Future()
_search_response.set_result({
'_shards': {
'total': 5,
'successful': 5
},
'_scroll_id': 42,
'hits': {
'hits': [1, 2, 3]
},
})
aes.search.return_value = _search_response
def _scroll_response(hits):
response = Future()
response.set_result({
'_shards': {
'total': 5,
'successful': 4
},
'_scroll_id': 42,
'hits': {
'hits': hits
},
})
return response
aes.scroll.side_effect = [
_scroll_response([4, 5, 6]),
_scroll_response([])
]
data = [
h async for h in scan(aes, raise_on_error=False, clear_scroll=False)
]
assert data == [1, 2, 3, 4, 5, 6]
assert aes.scroll.call_count == 2
aes.search.reset_mock()
aes.scroll.reset_mock()
aes.scroll.side_effect = [
_scroll_response([4, 5, 6]),
_scroll_response([])
]
with raises(ScanError):
data = []
async for h in scan(aes, raise_on_error=True, clear_scroll=False):
data.append(h)
assert data == [1, 2, 3]
assert aes.scroll.call_count == 1
async def test_clear_scroll(aes, index_name, test_data, mocker):
spy = mocker.patch.object(aes, 'clear_scroll', wraps=aes.clear_scroll)
async for _ in scan(aes, index=index_name):
pass
assert spy.call_count == 1
spy.reset_mock()
try:
async for _ in scan(aes, index=index_name, clear_scroll=False):
pass
assert spy.call_count == 0
finally:
await aes.transport.perform_request('DELETE', '/_search/scroll/_all')
async def test_fast_error_route():
aes = MagicMock(spec=AsyncElasticsearch)
async def _failed_response():
return {
'_shards': {
'total': 5,
'successful': 4
},
'_scroll_id': 42,
'hits': {
'hits': [1, 2, 3]
},
}
aes.search.return_value = _failed_response()
async def _clean_response():
return
aes.clear_scroll.return_value = _clean_response()
with raises(ScanError):
data = []
async for h in scan(aes):
data.append(h)
assert data == []
assert aes.scroll.call_count == 0
assert aes.clear_scroll.call_count == 1
async def test_initial_search_error():
aes = MagicMock(spec=AsyncElasticsearch)
async def _response():
return {}
aes.search.return_value = _response()
data = [h async for h in scan(aes)]
assert data == []
assert aes.scroll.call_count == 0
assert aes.clear_scroll.call_count == 0
async def scan(self):
"""
Turn the search into a scan search and return a async generator that will
iterate over all the documents matching the query.
Use ``params`` method to specify any additional arguments you with to pass to the
underlying ``scan`` helper
"""
aes = get_connection(self._using)
async for hit in scan(
aes,
query=self.to_dict(),
index=self._index,
**self._params
):
yield self._get_result(hit)
async def test_preserve_order(aes, index_name):
await aes.bulk(body=[
{
'index': {}
},
{
'value': 3
},
{
'index': {}
},
{
'value': 2
},
{
'index': {}
},
{
'value': 1
},
{
'index': {}
},
{
'value': 5
},
{
'index': {}
},
{
'value': 4
},
],
index=index_name,
refresh=True)
results = [
r async for r in scan(aes,
query={'sort': 'value'},
index=index_name,
preserve_order=True)
]
assert [r['_source']['value'] for r in results] == [1, 2, 3, 4, 5]
async def test_logger(mocker):
logger = mocker.patch('elasticsearch_adsl.scan.logger')
aes = MagicMock(spec=AsyncElasticsearch)
async def _failed_response():
return {
'_shards': {
'total': 5,
'successful': 4
},
'_scroll_id': 42,
'hits': {
'hits': [1, 2, 3]
},
}
aes.search.return_value = _failed_response()
with raises(ScanError):
async for _ in scan(aes, clear_scroll=False):
pass
assert logger.warning.call_count == 1
async def test_scan_basic(aes, test_data, index_name):
results = [r async for r in scan(aes, index=index_name)]
assert len(results) == 3
