def test_repr(self):
query = SearchQuery()
assert repr(query) == "<SearchQuery [None]>"
query = SearchQuery(search_text="foo bar", suggester_name="sg")
assert repr(query) == "<SearchQuery [foo bar]>"
query = SearchQuery(search_text="aaaaabbbbb" * 200)
assert len(repr(query)) == 1024
def test_get_search_facets_none(self, api_key, endpoint, index_name, **kwargs):
client = SearchIndexClient(
endpoint, index_name, AzureKeyCredential(api_key)
)
query = SearchQuery(search_text="WiFi")
query.select("hotelName", "category", "description")
results = client.search(query=query)
assert results.get_facets() is None
def test_get_search_counts(self, api_key, endpoint, index_name, **kwargs):
client = SearchIndexClient(
endpoint, index_name, AzureKeyCredential(api_key)
)
query = SearchQuery(search_text="hotel")
results = client.search(query=query)
assert results.get_count() is None
query = SearchQuery(search_text="hotel", include_total_result_count=True)
results = client.search(query=query)
assert results.get_count() == 7
def test_get_search_coverage(self, api_key, endpoint, index_name, **kwargs):
client = SearchIndexClient(
endpoint, index_name, AzureKeyCredential(api_key)
)
query = SearchQuery(search_text="hotel")
results = client.search(query=query)
assert results.get_coverage() is None
query = SearchQuery(search_text="hotel", minimum_coverage=50.0)
results = client.search(query=query)
cov = results.get_coverage()
assert isinstance(cov, float)
assert cov >= 50.0
async def test_get_search_filter(self, api_key, endpoint, index_name,
**kwargs):
client = SearchIndexClient(endpoint, index_name,
SearchApiKeyCredential(api_key))
query = SearchQuery(search_text="WiFi")
query.filter("category eq 'Budget'")
query.select("hotelName", "category", "description")
query.order_by("hotelName desc")
async with client:
results = []
async for x in await client.search(query=query):
results.append(x)
assert [x["hotelName"] for x in results
] == sorted([x["hotelName"] for x in results],
reverse=True)
expected = {
"category",
"hotelName",
"description",
"@search.score",
"@search.highlights",
}
assert all(set(x) == expected for x in results)
assert all(x["category"] == "Budget" for x in results)
def test_get_search_facets_result(self, api_key, endpoint, index_name, **kwargs):
client = SearchIndexClient(
endpoint, index_name, AzureKeyCredential(api_key)
)
query = SearchQuery(search_text="WiFi", facets=["category"])
query.select("hotelName", "category", "description")
results = client.search(query=query)
assert results.get_facets() == {
"category": [
{"value": "Budget", "count": 4},
{"value": "Luxury", "count": 1},
]
}
def test_filter(self):
query = SearchQuery()
assert query.request.filter is None
query.filter("expr0")
assert query.request.filter == "expr0"
query = SearchQuery(filter="expr0")
assert query.request.filter == "expr0"
query.filter("expr1")
assert query.request.filter == "expr1"
def filter_query():
# [START facet_query]
from azure.search.documents import SearchApiKeyCredential, SearchIndexClient, SearchQuery
search_client = SearchIndexClient(service_endpoint, index_name, SearchApiKeyCredential(key))
query = SearchQuery(search_text="WiFi", facets=["Category"], top=0)
results = search_client.search(query=query)
facets = results.get_facets()
print("Catgory facet counts for hotels:")
for facet in facets["Category"]:
print(" {}".format(facet))
def filter_query():
# [START filter_query]
from azure.search.documents import SearchApiKeyCredential, SearchIndexClient, SearchQuery
search_client = SearchIndexClient(service_endpoint, index_name,
SearchApiKeyCredential(key))
query = SearchQuery(search_text="WiFi")
query.filter("Address/StateProvince eq 'FL' and Address/Country eq 'USA'")
query.select("HotelName", "Rating")
query.order_by("Rating desc")
results = search_client.search(query=query)
print("Florida hotels containing 'WiFi', sorted by Rating:")
for result in results:
print(" Name: {} (rating {})".format(result["HotelName"],
result["Rating"]))
async def filter_query():
# [START facet_query_async]
from azure.core.credentials import AzureKeyCredential
from azure.search.documents.aio import SearchClient
from azure.search.documents import SearchQuery
search_client = SearchClient(service_endpoint, index_name,
AzureKeyCredential(key))
query = SearchQuery(search_text="WiFi", facets=["Category"], top=0)
async with search_client:
results = await search_client.search(query=query)
facets = await results.get_facets()
print("Catgory facet counts for hotels:")
for facet in facets["Category"]:
print(" {}".format(facet))
class TestSearchIndexClient(object):
def test_init(self):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
assert client._headers == {
"api-key": "test_api_key",
"Accept": "application/json;odata.metadata=none",
}
def test_credential_roll(self):
credential = SearchApiKeyCredential(api_key="old_api_key")
client = SearchIndexClient("endpoint", "index name", credential)
assert client._headers == {
"api-key": "old_api_key",
"Accept": "application/json;odata.metadata=none",
}
credential.update_key("new_api_key")
assert client._headers == {
"api-key": "new_api_key",
"Accept": "application/json;odata.metadata=none",
}
def test_headers_merge(self):
credential = SearchApiKeyCredential(api_key="test_api_key")
client = SearchIndexClient("endpoint", "index name", credential)
orig = {"foo": "bar"}
result = client._merge_client_headers(orig)
assert result is not orig
assert result == {
"api-key": "test_api_key",
"Accept": "application/json;odata.metadata=none",
"foo": "bar",
}
def test_repr(self):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
assert repr(
client) == "<SearchIndexClient [endpoint={}, index={}]>".format(
repr("endpoint"), repr("index name"))
@mock.patch(
"azure.search.documents._index._generated.operations._documents_operations.DocumentsOperations.count"
)
def test_get_document_count(self, mock_count):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
client.get_document_count()
assert mock_count.called
assert mock_count.call_args[0] == ()
assert len(mock_count.call_args[1]) == 1
assert mock_count.call_args[1]["headers"] == client._headers
@mock.patch(
"azure.search.documents._index._generated.operations._documents_operations.DocumentsOperations.get"
)
def test_get_document(self, mock_get):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
client.get_document("some_key")
assert mock_get.called
assert mock_get.call_args[0] == ()
assert len(mock_get.call_args[1]) == 3
assert mock_get.call_args[1]["headers"] == client._headers
assert mock_get.call_args[1]["key"] == "some_key"
assert mock_get.call_args[1]["selected_fields"] == None
mock_get.reset()
client.get_document("some_key", selected_fields="foo")
assert mock_get.called
assert mock_get.call_args[0] == ()
assert len(mock_get.call_args[1]) == 3
assert mock_get.call_args[1]["headers"] == client._headers
assert mock_get.call_args[1]["key"] == "some_key"
assert mock_get.call_args[1]["selected_fields"] == "foo"
@pytest.mark.parametrize(
"query",
["search text", SearchQuery(search_text="search text")],
ids=repr)
@mock.patch(
"azure.search.documents._index._generated.operations._documents_operations.DocumentsOperations.search_post"
)
def test_search_query_argument(self, mock_search_post, query):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
result = client.search(query)
assert isinstance(result, ItemPaged)
assert result._page_iterator_class is SearchPageIterator
search_result = SearchDocumentsResult()
search_result.results = [
SearchResult(additional_properties={"key": "val"})
]
mock_search_post.return_value = search_result
assert not mock_search_post.called
next(result)
assert mock_search_post.called
assert mock_search_post.call_args[0] == ()
assert (mock_search_post.call_args[1]["search_request"].search_text ==
"search text")
def test_search_bad_argument(self):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
with pytest.raises(TypeError) as e:
client.search(10)
assert str(
e) == "Expected a SuggestQuery for 'query', but got {}".format(
repr(10))
@mock.patch(
"azure.search.documents._index._generated.operations._documents_operations.DocumentsOperations.suggest_post"
)
def test_suggest_query_argument(self, mock_suggest_post):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
result = client.suggest(
SuggestQuery(search_text="search text", suggester_name="sg"))
assert mock_suggest_post.called
assert mock_suggest_post.call_args[0] == ()
assert mock_suggest_post.call_args[1]["headers"] == client._headers
assert (mock_suggest_post.call_args[1]["suggest_request"].search_text
== "search text")
def test_suggest_bad_argument(self):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
with pytest.raises(TypeError) as e:
client.suggest("bad_query")
assert str(
e) == "Expected a SuggestQuery for 'query', but got {}".format(
repr("bad_query"))
@mock.patch(
"azure.search.documents._index._generated.operations._documents_operations.DocumentsOperations.autocomplete_post"
)
def test_autocomplete_query_argument(self, mock_autocomplete_post):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
result = client.autocomplete(
AutocompleteQuery(search_text="search text", suggester_name="sg"))
assert mock_autocomplete_post.called
assert mock_autocomplete_post.call_args[0] == ()
assert mock_autocomplete_post.call_args[1][
"headers"] == client._headers
assert (mock_autocomplete_post.call_args[1]
["autocomplete_request"].search_text == "search text")
def test_autocomplete_bad_argument(self):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
with pytest.raises(TypeError) as e:
client.autocomplete("bad_query")
assert str(
e
) == "Expected a AutocompleteQuery for 'query', but got {}".format(
repr("bad_query"))
@pytest.mark.parametrize("arg", [[], ["doc1"], ["doc1", "doc2"]],
ids=lambda x: str(len(x)) + " docs")
@pytest.mark.parametrize("method_name", CRUD_METHOD_NAMES)
def test_add_method(self, arg, method_name):
with mock.patch.object(SearchIndexClient,
"index_documents",
return_value=None) as mock_index_documents:
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
method = getattr(client, method_name)
method(arg, extra="foo")
assert mock_index_documents.called
assert len(mock_index_documents.call_args[0]) == 1
batch = mock_index_documents.call_args[0][0]
assert isinstance(batch, IndexDocumentsBatch)
assert all(action.action_type == CRUD_METHOD_MAP[method_name]
for action in batch.actions)
assert [action.additional_properties
for action in batch.actions] == arg
assert mock_index_documents.call_args[1][
"headers"] == client._headers
assert mock_index_documents.call_args[1]["extra"] == "foo"
@mock.patch(
"azure.search.documents._index._generated.operations._documents_operations.DocumentsOperations.index"
)
def test_index_documents(self, mock_index):
client = SearchIndexClient("endpoint", "index name", CREDENTIAL)
batch = IndexDocumentsBatch()
batch.add_upload_documents("upload1")
batch.add_delete_documents("delete1", "delete2")
batch.add_merge_documents(["merge1", "merge2", "merge3"])
batch.add_merge_or_upload_documents("merge_or_upload1")
client.index_documents(batch, extra="foo")
assert mock_index.called
assert mock_index.call_args[0] == ()
assert len(mock_index.call_args[1]) == 3
assert mock_index.call_args[1]["headers"] == client._headers
assert mock_index.call_args[1]["extra"] == "foo"
index_documents = mock_index.call_args[1]["batch"]
assert isinstance(index_documents, IndexBatch)
assert index_documents.actions == batch.actions
def test_order_by(self):
query = SearchQuery()
assert query.request.order_by is None
query.order_by("f0")
assert query.request.order_by == "f0"
query.order_by("f1,f2")
assert query.request.order_by == "f1,f2"
query.order_by("f3", "f4")
assert query.request.order_by == "f3,f4"
query = SearchQuery(order_by="f0")
assert query.request.order_by == "f0"
query.order_by("f1,f2")
assert query.request.order_by == "f1,f2"
query.order_by("f3", "f4")
assert query.request.order_by == "f3,f4"
with pytest.raises(ValueError) as e:
query.order_by()
assert str(e) == "At least one field must be provided"
def test_kwargs_forwarded(self, mock_request):
mock_request.return_value = None
SearchQuery(foo=10, bar=20)
assert mock_request.called
assert mock_request.call_args[0] == ()
assert mock_request.call_args[1] == {"foo": 10, "bar": 20}
def test_init(self):
query = SearchQuery()
assert type(query.request) is SearchRequest
assert query.request.filter is None
assert query.request.order_by is None
assert query.request.select is None
def test_select(self):
query = SearchQuery()
assert query.request.select is None
query.select("f0")
assert query.request.select == "f0"
query.select("f1,f2")
assert query.request.select == "f1,f2"
query.select("f3", "f4")
assert query.request.select == "f3,f4"
query = SearchQuery(select="f0")
assert query.request.select == "f0"
query.select("f1,f2")
assert query.request.select == "f1,f2"
query.select("f3", "f4")
assert query.request.select == "f3,f4"
with pytest.raises(ValueError) as e:
query.select()
assert str(e) == "At least one field must be provided"
