search perf tests

sdk/search/azure-search-documents/tests/perfstress_tests/README.md

@@ -0,0 +1,47 @@
+# Search Performance Tests
+
+In order to run the performance tests, the `azure-devtools` package must be installed. This is done as part of the `dev_requirements`.
+Start be creating a new virtual environment for your perf tests. This will need to be a Python 3 environment, preferably >=3.7.
+
+### Setup for test resources
+
+These tests will run against a pre-configured search service. See [here](https://docs.microsoft.com/azure/search/search-indexer-tutorial) about how to configure the service and import data. The following environment variable will need to be set for the tests to access the live resources:
+```
+AZURE_SEARCH_API_KEY=<search api key>
+AZURE_SEARCH_SERVICE_ENDPOINT=<end point url>
+AZURE_SEARCH_INDEX_NAME=<index name>
+```
+
+### Setup for perf test runs
+
+```cmd
+(env) ~/search/azure-search-documents> pip install -r dev_requirements.txt
+(env) ~/search/azure-search-documents> pip install -e .
+```
+
+## Test commands
+
+When `azure-devtools` is installed, you will have access to the `perfstress` command line tool, which will scan the current module for runable perf tests. Only a specific test can be run at a time (i.e. there is no "run all" feature).
+
+```cmd
+(env) ~/search/azure-search-documents> cd tests
+(env) ~/search/azure-search-documents/tests> perfstress
+```
+Using the `perfstress` command alone will list the available perf tests found. 
+
+### Common perf command line options
+These options are available for all perf tests:
+- `--duration=10` Number of seconds to run as many operations (the "run" function) as possible. Default is 10.
+- `--iterations=1` Number of test iterations to run. Default is 1.
+- `--parallel=1` Number of tests to run in parallel. Default is 1.
+- `--warm-up=5` Number of seconds to spend warming up the connection before measuring begins. Default is 5.
+- `--sync` Whether to run the tests in sync or async. Default is False (async). This flag must be used for Storage legacy tests, which do not support async.
+- `--no-cleanup` Whether to keep newly created resources after test run. Default is False (resources will be deleted).
+- `--num-documents` The number of results expect to be returned.
+
+## Example command
+```cmd
+(env) ~/search/azure-search-documents/tests> perfstress SearchDocumentsTest
+(env) ~/search/azure-search-documents/tests> perfstress AutoCompleteTest
+(env) ~/search/azure-search-documents/tests> perfstress SuggestTest
+```

sdk/search/azure-search-documents/tests/perfstress_tests/autocomplete.py

@@ -0,0 +1,46 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import os
+from azure_devtools.perfstress_tests import PerfStressTest
+
+from azure.core.credentials import AzureKeyCredential
+from azure.search.documents import SearchClient as SyncClient
+from azure.search.documents.aio import SearchClient as AsyncClient
+
+
+class AutoCompleteTest(PerfStressTest):
+    def __init__(self, arguments):
+        super().__init__(arguments)
+        api_key = self.get_from_env("AZURE_SEARCH_API_KEY")
+        service_endpoint = os.getenv("AZURE_SEARCH_SERVICE_ENDPOINT")
+        index_name = os.getenv("AZURE_SEARCH_INDEX_NAME")
+        key = os.getenv("AZURE_SEARCH_API_KEY")
+        self.service_client = SyncClient(service_endpoint, index_name, AzureKeyCredential(api_key))
+        self.async_service_client = AsyncClient(service_endpoint, index_name, AzureKeyCredential(api_key))
+
+    async def close(self):
+        await self.async_service_client.close()
+        await super().close()
+
+    @staticmethod
+    def add_arguments(parser):
+        super(AutoCompleteTest, AutoCompleteTest).add_arguments(parser)
+        parser.add_argument('--num-documents', nargs='?', type=int,
+                            help='The number of results expect to be returned.',
+                            default=-1)
+
+    def run_sync(self):
+        if self.args.num_documents == -1:
+            results = len(self.service_client.autocomplete(search_text="mot", suggester_name="sg"))
+        else:
+            results = len(self.service_client.autocomplete(search_text="mot", suggester_name="sg", top=self.args.num_documents))
+
+    async def run_async(self):
+        if self.args.num_documents == -1:
+            results = len(await self.async_service_client.autocomplete(search_text="mot", suggester_name="sg"))
+        else:
+            results = len(await self.async_service_client.autocomplete(search_text="mot", suggester_name="sg", top=self.args.num_documents))
+

sdk/search/azure-search-documents/tests/perfstress_tests/search_documents.py

@@ -0,0 +1,53 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import os
+
+from azure_devtools.perfstress_tests import PerfStressTest
+
+from azure.core.credentials import AzureKeyCredential
+from azure.search.documents import SearchClient as SyncClient
+from azure.search.documents.aio import SearchClient as AsyncClient
+
+
+class SearchDocumentsTest(PerfStressTest):
+    def __init__(self, arguments):
+        super().__init__(arguments)
+        api_key = self.get_from_env("AZURE_SEARCH_API_KEY")
+        service_endpoint = os.getenv("AZURE_SEARCH_SERVICE_ENDPOINT")
+        index_name = os.getenv("AZURE_SEARCH_INDEX_NAME")
+        key = os.getenv("AZURE_SEARCH_API_KEY")
+        self.service_client = SyncClient(service_endpoint, index_name, AzureKeyCredential(api_key))
+        self.async_service_client = AsyncClient(service_endpoint, index_name, AzureKeyCredential(api_key))
+
+    @staticmethod
+    def add_arguments(parser):
+        super(SearchDocumentsTest, SearchDocumentsTest).add_arguments(parser)
+        parser.add_argument('--num-documents', nargs='?', type=int,
+                            help='The number of results expect to be returned.',
+                            default=-1)
+
+    async def global_setup(self):
+        await super().global_setup()
+
+    async def close(self):
+        await self.async_service_client.close()
+        await super().close()
+
+    def run_sync(self):
+        if self.args.num_documents == -1:
+            results = len(self.service_client.search(search_text="luxury"))
+        else:
+            results = len(self.service_client.search(search_text="luxury", top=self.args.num_documents))
+
+
+    async def run_async(self):
+        if self.args.num_documents == -1:
+            results = await self.async_service_client.search(search_text="luxury")
+        else:
+            results = await self.async_service_client.search(search_text="luxury", top=self.args.num_documents)
+        count = 0
+        async for result in results:
+            count += count

sdk/search/azure-search-documents/tests/perfstress_tests/suggest.py

@@ -0,0 +1,46 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import os
+from azure_devtools.perfstress_tests import PerfStressTest
+
+from azure.core.credentials import AzureKeyCredential
+from azure.search.documents import SearchClient as SyncClient
+from azure.search.documents.aio import SearchClient as AsyncClient
+
+
+class SuggestTest(PerfStressTest):
+    def __init__(self, arguments):
+        super().__init__(arguments)
+        api_key = self.get_from_env("AZURE_SEARCH_API_KEY")
+        service_endpoint = os.getenv("AZURE_SEARCH_SERVICE_ENDPOINT")
+        index_name = os.getenv("AZURE_SEARCH_INDEX_NAME")
+        key = os.getenv("AZURE_SEARCH_API_KEY")
+        self.service_client = SyncClient(service_endpoint, index_name, AzureKeyCredential(api_key))
+        self.async_service_client = AsyncClient(service_endpoint, index_name, AzureKeyCredential(api_key))
+
+    async def close(self):
+        await self.async_service_client.close()
+        await super().close()
+
+    @staticmethod
+    def add_arguments(parser):
+        super(SuggestTest, SuggestTest).add_arguments(parser)
+        parser.add_argument('--num-documents', nargs='?', type=int,
+                            help='The number of results expect to be returned.',
+                            default=-1)
+
+    def run_sync(self):
+        if self.args.num_documents == -1:
+            results = len(self.service_client.suggest(search_text="mot", suggester_name="sg"))
+        else:
+            results = len(self.service_client.suggest(search_text="mot", suggester_name="sg", top=self.args.num_documents))
+
+    async def run_async(self):
+        if self.args.num_documents == -1:
+            results = len(await self.async_service_client.suggest(search_text="mot", suggester_name="sg"))
+        else:
+            results = len(await self.async_service_client.suggest(search_text="mot", suggester_name="sg", top=self.args.num_documents))
+