async def test_upload_documents_existing(self, api_key, endpoint,
index_name, **kwargs):
client = SearchClient(endpoint, index_name,
AzureKeyCredential(api_key))
batch_client = SearchIndexingBufferedSender(
endpoint, index_name, AzureKeyCredential(api_key))
batch_client._batch_action_count = 2
DOCUMENTS = [
{
"hotelId": "1000",
"rating": 5,
"rooms": [],
"hotelName": "Azure Inn"
},
{
"hotelId": "3",
"rating": 4,
"rooms": [],
"hotelName": "Redmond Hotel"
},
]
async with batch_client:
await batch_client.upload_documents(DOCUMENTS)
# There can be some lag before a document is searchable
if self.is_live:
time.sleep(TIME_TO_SLEEP)
async with client:
assert await client.get_document_count() == 11
async def test_delete_documents_missing(self, api_key, endpoint,
index_name, **kwargs):
client = SearchClient(endpoint, index_name,
AzureKeyCredential(api_key))
batch_client = SearchIndexingBufferedSender(
endpoint, index_name, AzureKeyCredential(api_key))
batch_client._batch_action_count = 2
async with batch_client:
await batch_client.delete_documents([{
"hotelId": "1000"
}, {
"hotelId": "4"
}])
# There can be some lag before a document is searchable
if self.is_live:
time.sleep(TIME_TO_SLEEP)
async with client:
assert await client.get_document_count() == 9
with pytest.raises(HttpResponseError):
await client.get_document(key="1000")
with pytest.raises(HttpResponseError):
await client.get_document(key="4")
async def test_merge_or_upload_documents(self, api_key, endpoint,
index_name, **kwargs):
client = SearchClient(endpoint, index_name,
AzureKeyCredential(api_key))
batch_client = SearchIndexingBufferedSender(
endpoint, index_name, AzureKeyCredential(api_key))
batch_client._batch_action_count = 2
async with batch_client:
await batch_client.merge_or_upload_documents([{
"hotelId": "1000",
"rating": 1
}, {
"hotelId": "4",
"rating": 2
}])
# There can be some lag before a document is searchable
if self.is_live:
time.sleep(TIME_TO_SLEEP)
async with client:
assert await client.get_document_count() == 11
result = await client.get_document(key="1000")
assert result["rating"] == 1
result = await client.get_document(key="4")
assert result["rating"] == 2
async def test_search_client_index_buffered_sender(self, endpoint, api_key,
index_name):
client = SearchClient(endpoint, index_name, api_key)
batch_client = SearchIndexingBufferedSender(endpoint, index_name,
api_key)
try:
async with client:
async with batch_client:
doc_count = 10
doc_count = await self._test_upload_documents_new(
client, batch_client, doc_count)
doc_count = await self._test_upload_documents_existing(
client, batch_client, doc_count)
doc_count = await self._test_delete_documents_existing(
client, batch_client, doc_count)
doc_count = await self._test_delete_documents_missing(
client, batch_client, doc_count)
doc_count = await self._test_merge_documents_existing(
client, batch_client, doc_count)
doc_count = await self._test_merge_documents_missing(
client, batch_client, doc_count)
doc_count = await self._test_merge_or_upload_documents(
client, batch_client, doc_count)
finally:
batch_client.close()
async def test_context_manager(self, mock_cleanup):
async with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
auto_flush=False) as client:
await client.upload_documents(["upload1"])
await client.delete_documents(["delete1", "delete2"])
assert mock_cleanup.called
async def test_callback_new(self):
on_new = mock.AsyncMock()
async with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
auto_flush=False,
on_new=on_new) as client:
await client.upload_documents(["upload1"])
assert on_new.called
async def test_search_indexing_buffered_sender_kwargs(self):
async with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
window=100) as client:
assert client._batch_action_count == 512
assert client._max_retries == 3
assert client._auto_flush_interval == 60
assert client._auto_flush
async def test_batch_queue(self):
async with SearchIndexingBufferedSender("endpoint", "index name", CREDENTIAL, auto_flush=False) as client:
assert client._index_documents_batch
await client.upload_documents(["upload1"])
await client.delete_documents(["delete1", "delete2"])
await client.merge_documents(["merge1", "merge2", "merge3"])
await client.merge_or_upload_documents(["merge_or_upload1"])
assert len(client.actions) == 7
actions = await client._index_documents_batch.dequeue_actions()
assert len(client.actions) == 0
await client._index_documents_batch.enqueue_actions(actions)
assert len(client.actions) == 7
async def test_flush(self):
DOCUMENT = {
'Category': 'Hotel',
'HotelId': '1000',
'Rating': 4.0,
'Rooms': [],
'HotelName': 'Azure Inn',
}
with mock.patch.object(SearchIndexingBufferedSender, "_index_documents_actions", side_effect=HttpResponseError("Error")):
async with SearchIndexingBufferedSender("endpoint", "index name", CREDENTIAL, auto_flush=False) as client:
client._index_key = "HotelId"
await client.upload_documents([DOCUMENT])
await client.flush()
assert len(client.actions) == 0
async def test_callback_error(self):
async def mock_fail_index_documents(actions, timeout=86400):
if len(actions) > 0:
result = IndexingResult()
result.key = actions[0].additional_properties.get('id')
result.status_code = 400
result.succeeded = False
self.uploaded = self.uploaded + len(actions) - 1
return [result]
on_error = mock.AsyncMock()
async with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
auto_flush=False,
on_error=on_error) as client:
client._index_documents_actions = mock_fail_index_documents
client._index_key = "id"
await client.upload_documents({"id": 0})
await client.flush()
assert on_error.called
async def sample_batching_client():
DOCUMENT = {
'Category': 'Hotel',
'HotelId': '1000',
'Rating': 4.0,
'Rooms': [],
'HotelName': 'Azure Inn',
}
async with SearchIndexingBufferedSender(
service_endpoint, index_name,
AzureKeyCredential(key)) as batch_client:
# add upload actions
await batch_client.upload_documents(documents=[DOCUMENT])
# add merge actions
await batch_client.merge_documents(documents=[{
"HotelId": "1000",
"Rating": 4.5
}])
# add delete actions
await batch_client.delete_documents(documents=[{"HotelId": "1000"}])
async def test_callback_progress(self):
async def mock_successful_index_documents(actions, timeout=86400):
if len(actions) > 0:
result = IndexingResult()
result.key = actions[0].additional_properties.get('id')
result.status_code = 200
result.succeeded = True
return [result]
on_progress = mock.AsyncMock()
on_remove = mock.AsyncMock()
async with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
auto_flush=False,
on_progress=on_progress,
on_remove=on_remove) as client:
client._index_documents_actions = mock_successful_index_documents
client._index_key = "id"
await client.upload_documents({"id": 0})
await client.flush()
assert on_progress.called
assert on_remove.called
async def test_callback_error_on_timeout(self):
async def mock_fail_index_documents(actions, timeout=86400):
if len(actions) > 0:
result = IndexingResult()
result.key = actions[0].additional_properties.get('id')
result.status_code = 400
result.succeeded = False
self.uploaded = self.uploaded + len(actions) - 1
time.sleep(1)
return [result]
on_error = mock.AsyncMock()
async with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
auto_flush=False,
on_error=on_error) as client:
client._index_documents_actions = mock_fail_index_documents
client._index_key = "id"
await client.upload_documents([{"id": 0}, {"id": 1}])
with pytest.raises(ServiceResponseTimeoutError):
await client.flush(timeout=-1)
assert on_error.call_count == 2
async def test_process_if_needed(self, mock_process_if_needed):
async with SearchIndexingBufferedSender("endpoint", "index name",
CREDENTIAL) as client:
await client.upload_documents(["upload1"])
await client.delete_documents(["delete1", "delete2"])
assert mock_process_if_needed.called