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_size = 2
DOCUMENTS = [
{
"hotelId": "1000",
"rating": 5,
"rooms": [],
"hotelName": "Azure Inn"
},
{
"hotelId": "3",
"rating": 4,
"rooms": [],
"hotelName": "Redmond Hotel"
},
]
batch_client.upload_documents(DOCUMENTS)
# There can be some lag before a document is searchable
if self.is_live:
time.sleep(TIME_TO_SLEEP)
assert client.get_document_count() == 11
batch_client.close()
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_size = 2
batch_client.merge_or_upload_documents([{
"hotelId": "1000",
"rating": 1
}, {
"hotelId": "4",
"rating": 2
}])
batch_client.close()
# There can be some lag before a document is searchable
if self.is_live:
time.sleep(TIME_TO_SLEEP)
assert client.get_document_count() == 11
result = client.get_document(key="1000")
assert result["rating"] == 1
result = client.get_document(key="4")
assert result["rating"] == 2
def test_context_manager(self, mock_cleanup):
with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
auto_flush=False) as client:
client.upload_documents(["upload1"])
client.delete_documents(["delete1", "delete2"])
assert mock_cleanup.called
def test_callback_new(self):
on_new = mock.Mock()
with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
auto_flush=False,
on_new=on_new) as client:
client.upload_documents(["upload1"])
assert on_new.called
def test_search_indexing_buffered_sender_kwargs(self):
with SearchIndexingBufferedSender("endpoint",
"index name",
CREDENTIAL,
window=100) as client:
assert client._batch_action_count == 512
assert client._max_retries_per_action == 3
assert client._auto_flush_interval == 60
assert client._auto_flush
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")):
with SearchIndexingBufferedSender("endpoint", "index name", CREDENTIAL, auto_flush=False) as client:
client._index_key = "HotelId"
client.upload_documents([DOCUMENT])
client.flush()
assert len(client.actions) == 0
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:
doc_count = 10
doc_count = self._test_upload_documents_new(client, batch_client, doc_count)
doc_count = self._test_upload_documents_existing(client, batch_client, doc_count)
doc_count = self._test_delete_documents_existing(client, batch_client, doc_count)
doc_count = self._test_delete_documents_missing(client, batch_client, doc_count)
doc_count = self._test_merge_documents_existing(client, batch_client, doc_count)
doc_count = self._test_merge_documents_missing(client, batch_client, doc_count)
doc_count = self._test_merge_or_upload_documents(client, batch_client, doc_count)
finally:
batch_client.close()
def test_batch_queue(self):
with SearchIndexingBufferedSender("endpoint", "index name", CREDENTIAL, auto_flush=False) as client:
assert client._index_documents_batch
client.upload_documents(["upload1"])
client.delete_documents(["delete1", "delete2"])
client.merge_documents(["merge1", "merge2", "merge3"])
client.merge_or_upload_documents(["merge_or_upload1"])
assert len(client.actions) == 7
actions = client._index_documents_batch.dequeue_actions()
assert len(client.actions) == 0
client._index_documents_batch.enqueue_actions(actions)
assert len(client.actions) == 7
actions = client._index_documents_batch.dequeue_actions()
assert len(client.actions) == 0
for action in actions:
client._index_documents_batch.enqueue_action(action)
assert len(client.actions) == 7
def sample_batching_client():
DOCUMENT = {
'Category': 'Hotel',
'HotelId': '1000',
'Rating': 4.0,
'Rooms': [],
'HotelName': 'Azure Inn',
}
with SearchIndexingBufferedSender(
service_endpoint,
index_name,
AzureKeyCredential(key)) as batch_client:
# add upload actions
batch_client.upload_documents(documents=[DOCUMENT])
# add merge actions
batch_client.merge_documents(documents=[{"HotelId": "1000", "Rating": 4.5}])
# add delete actions
batch_client.delete_documents(documents=[{"HotelId": "1000"}])
def test_callback_error(self):
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.Mock()
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"
client.upload_documents({"id": 0})
client.flush()
assert on_error.called
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_size = 2
batch_client.delete_documents([{"hotelId": "1000"}, {"hotelId": "4"}])
batch_client.close()
# There can be some lag before a document is searchable
if self.is_live:
time.sleep(TIME_TO_SLEEP)
assert client.get_document_count() == 9
with pytest.raises(HttpResponseError):
client.get_document(key="1000")
with pytest.raises(HttpResponseError):
client.get_document(key="4")
def test_callback_progress(self):
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.Mock()
on_remove = mock.Mock()
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"
client.upload_documents({"id": 0})
client.flush()
assert on_progress.called
assert on_remove.called
def test_callback_error_on_timeout(self):
def mock_fail_index_documents(actions, timeout=86400):
import time
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.Mock()
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"
client.upload_documents([{"id": 0}, {"id": 1}])
with pytest.raises(ServiceResponseTimeoutError):
client.flush(timeout=-1)
assert on_error.call_count == 2
def test_merge_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
batch_client.merge_documents(
[{"hotelId": "1000", "rating": 1}, {"hotelId": "4", "rating": 2}]
)
batch_client.close()
# There can be some lag before a document is searchable
if self.is_live:
time.sleep(TIME_TO_SLEEP)
assert client.get_document_count() == 10
with pytest.raises(HttpResponseError):
client.get_document(key="1000")
result = client.get_document(key="4")
assert result["rating"] == 2
def test_process_if_needed(self, mock_process_if_needed):
with SearchIndexingBufferedSender("endpoint", "index name",
CREDENTIAL) as client:
client.upload_documents(["upload1"])
client.delete_documents(["delete1", "delete2"])
assert mock_process_if_needed.called