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