More test fixes

Author: TimPansino

tests/datastore_elasticsearch/test_async_elasticsearch.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import elasticsearch._async.client as client
-import pytest
 from conftest import ES_SETTINGS, ES_VERSION
 from testing_support.fixture.event_loop import event_loop as loop  # noqa: F401
 from testing_support.fixtures import override_application_settings

tests/datastore_elasticsearch/test_async_instrumented_methods.py

@@ -66,9 +66,9 @@ def test_method_on_async_client_datastore_trace_inputs(
     @background_task()
     async def _test():
         if not sub_module:
-            getattr(async_client, method)(*args, **kwargs)
+            await getattr(async_client, method)(*args, **kwargs)
         else:
-            getattr(getattr(async_client, sub_module), method)(*args, **kwargs)
+            await getattr(getattr(async_client, sub_module), method)(*args, **kwargs)
 
     loop.run_until_complete(_test())
 

tests/datastore_elasticsearch/test_async_mget.py

@@ -73,17 +73,19 @@
 
 
 @pytest.fixture(scope="module")
-def client():
+def client(loop):
     urls = [f"http://{db['host']}:{db['port']}" for db in ES_MULTIPLE_SETTINGS]
     # When selecting a connection from the pool, use the round robin method.
     # This is actually the default already. Using round robin will ensure that
     # doing two db calls will mean elastic search is talking to two different
     # dbs.
     if ES_VERSION >= (8,):
-        client = AsyncElasticsearch(urls, node_selector_class=RoundRobinSelector, randomize_hosts=False)
+        _client = AsyncElasticsearch(urls, node_selector_class=RoundRobinSelector, randomize_nodes_in_pool=False)
     else:
-        client = AsyncElasticsearch(urls, selector_class=RoundRobinSelector, randomize_hosts=False)
-    return client
+        _client = AsyncElasticsearch(urls, selector_class=RoundRobinSelector, randomize_nodes_in_pool=False)
+
+    yield _client
+    loop.run_until_complete(_client.close())
 
 
 # Query

tests/datastore_elasticsearch/test_async_multiple_dbs.py

@@ -78,6 +78,19 @@ async def _exercise_es(es):
 # Test
 
 
+@pytest.fixture(scope="session")
+def async_clients(loop):
+    clients = []
+    for db in ES_MULTIPLE_SETTINGS:
+        es_url = f"http://{db['host']}:{db['port']}"
+        clients.append(AsyncElasticsearch(es_url))
+
+    yield clients
+
+    for client in clients:
+        loop.run_until_complete(client.close())
+
+
 @pytest.mark.skipif(len(ES_MULTIPLE_SETTINGS) < 2, reason="Test environment not configured with multiple databases.")
 @override_application_settings(_enable_instance_settings)
 @validate_transaction_metrics(
@@ -87,11 +100,11 @@ async def _exercise_es(es):
     background_task=True,
 )
 @background_task()
-def test_async_multiple_dbs_enabled(loop):
-    for db in ES_MULTIPLE_SETTINGS:
-        es_url = f"http://{db['host']}:{db['port']}"
-        client = AsyncElasticsearch(es_url)
-        loop.run_until_complete(_exercise_es(client))
+def test_async_multiple_dbs_enabled(loop, async_clients):
+    import asyncio
+
+    # Run multiple queries in parallel
+    loop.run_until_complete(asyncio.gather(*(_exercise_es(client) for client in async_clients)))
 
 
 @pytest.mark.skipif(len(ES_MULTIPLE_SETTINGS) < 2, reason="Test environment not configured with multiple databases.")
@@ -103,8 +116,8 @@ def test_async_multiple_dbs_enabled(loop):
     background_task=True,
 )
 @background_task()
-def test_async_multiple_dbs_disabled(loop):
-    for db in ES_MULTIPLE_SETTINGS:
-        es_url = f"http://{db['host']}:{db['port']}"
-        client = AsyncElasticsearch(es_url)
-        loop.run_until_complete(_exercise_es(client))
+def test_async_multiple_dbs_disabled(loop, async_clients):
+    import asyncio
+
+    # Run multiple queries in parallel
+    loop.run_until_complete(asyncio.gather(*(_exercise_es(client) for client in async_clients)))

tests/datastore_elasticsearch/test_mget.py

@@ -78,10 +78,12 @@ def client():
     # doing two db calls will mean elastic search is talking to two different
     # dbs.
     if ES_VERSION >= (8,):
-        client = Elasticsearch(urls, node_selector_class=RoundRobinSelector, randomize_hosts=False)
+        _client = Elasticsearch(urls, node_selector_class=RoundRobinSelector, randomize_nodes_in_pool=False)
     else:
-        client = Elasticsearch(urls, selector_class=RoundRobinSelector, randomize_hosts=False)
-    return client
+        _client = Elasticsearch(urls, selector_class=RoundRobinSelector, randomize_nodes_in_pool=False)
+
+    yield _client
+    _client.close()
 
 
 # Query