async def test_upload_documents_existing(self, api_key, endpoint,
index_name, **kwargs):
client = SearchClient(endpoint, index_name,
AzureKeyCredential(api_key))
batch_client = SearchIndexDocumentBatchingClient(
endpoint, index_name, AzureKeyCredential(api_key))
batch_client._batch_size = 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.add_upload_actions(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 = SearchIndexDocumentBatchingClient(
endpoint, index_name, AzureKeyCredential(api_key))
batch_client._batch_size = 2
async with batch_client:
await batch_client.add_delete_actions([{
"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 = SearchIndexDocumentBatchingClient(
endpoint, index_name, AzureKeyCredential(api_key))
batch_client._batch_size = 2
async with batch_client:
await batch_client.add_merge_or_upload_actions([{
"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_index_document_batching_client_kwargs(self):
client = SearchIndexDocumentBatchingClient("endpoint", "index name", CREDENTIAL, window=100)
assert client.batch_size == 1000
assert client._window == 100
assert client._auto_flush
await client.close()
async def test_flush_if_needed(self, mock_flush):
client = SearchIndexDocumentBatchingClient("endpoint", "index name", CREDENTIAL, window=1000, batch_size=2)
await client.add_upload_actions(["upload1"])
await client.add_delete_actions(["delete1", "delete2"])
assert mock_flush.called
await client.close()
async def test_failed_queue(self):
client = SearchIndexDocumentBatchingClient("endpoint", "index name", CREDENTIAL)
assert client._index_documents_batch
await client.add_upload_actions(["upload1"])
await client.add_delete_actions(["delete1", "delete2"])
await client.add_merge_actions(["merge1", "merge2", "merge3"])
await client.add_merge_or_upload_actions(["merge_or_upload1"])
actions = await client._index_documents_batch.dequeue_actions()
await client._index_documents_batch.enqueue_failed_actions(actions)
assert len(client.failed_actions) == 7
async def test_flush(self):
DOCUMENT = {
'Category': 'Hotel',
'HotelId': '1000',
'Rating': 4.0,
'Rooms': [],
'HotelName': 'Azure Inn',
}
with mock.patch.object(SearchIndexDocumentBatchingClient, "_index_documents_actions", side_effect=HttpResponseError("Error")):
async with SearchIndexDocumentBatchingClient("endpoint", "index name", CREDENTIAL, auto_flush=False) as client:
client._index_key = "HotelId"
await client.add_upload_actions([DOCUMENT])
await client.flush()
assert len(client.actions) == 0
async def sample_batching_client():
DOCUMENT = {
'Category': 'Hotel',
'HotelId': '1000',
'Rating': 4.0,
'Rooms': [],
'HotelName': 'Azure Inn',
}
async with SearchIndexDocumentBatchingClient(
service_endpoint,
index_name,
AzureKeyCredential(key),
window=100,
batch_size=100) as batch_client:
# add upload actions
await batch_client.add_upload_actions(documents=[DOCUMENT])
# add merge actions
await batch_client.add_merge_actions(documents=[{
"HotelId": "1000",
"Rating": 4.5
}])
# add delete actions
await batch_client.add_delete_actions(documents=[{"HotelId": "1000"}])
async def test_context_manager(self, mock_cleanup):
async with SearchIndexDocumentBatchingClient("endpoint", "index name", CREDENTIAL, auto_flush=False) as client:
await client.add_upload_actions(["upload1"])
await client.add_delete_actions(["delete1", "delete2"])
assert mock_cleanup.called
async def test_process_if_needed(self, mock_process_if_needed):
client = SearchIndexDocumentBatchingClient("endpoint", "index name", CREDENTIAL, window=1000, auto_flush=False)
await client.add_upload_actions(["upload1"])
await client.add_delete_actions(["delete1", "delete2"])
assert mock_process_if_needed.called
