add tests with async grpc server

Author: felipou

tests/contrib/grpc/grpc_app/server.py

@@ -28,7 +28,9 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import asyncio
 import logging
+import os
 import sys
 from concurrent import futures
 
@@ -67,6 +69,30 @@ def GetServerResponseException(self, request, context):
  raise Exception("oh no")
 
 
+class TestServiceAsync(pb2_grpc.TestServiceServicer):
+ def __init__(self, *args, **kwargs):
+ pass
+
+ async def GetServerResponse(self, request, context):
+ message = request.message
+ result = f'Hello I am up and running received "{message}" message from you'
+ result = {"message": result, "received": True}
+
+ return pb2.MessageResponse(**result)
+
+ async def GetServerResponseAbort(self, request, context):
+ await context.abort(grpc.StatusCode.INTERNAL, "foo")
+
+ async def GetServerResponseUnavailable(self, request, context):
+ """Missing associated documentation comment in .proto file."""
+ context.set_code(grpc.StatusCode.UNAVAILABLE)
+ context.set_details("Method not available")
+ return pb2.MessageResponse(message="foo", received=True)
+
+ async def GetServerResponseException(self, request, context):
+ raise Exception("oh no")
+
+
 def serve(port):
  apm_client = GRPCApmClient(
  service_name="grpc-server", disable_metrics="*", api_request_time="100ms", central_config="False"
@@ -78,10 +104,24 @@ def serve(port):
  server.wait_for_termination()
 
 
+async def serve_async(port):
+ apm_client = GRPCApmClient(
+ service_name="grpc-server", disable_metrics="*", api_request_time="100ms", central_config="False"
+ )
+ server = grpc.aio.server()
+ pb2_grpc.add_TestServiceServicer_to_server(TestServiceAsync(), server)
+ server.add_insecure_port(f"[::]:{port}")
+ await server.start()
+ await server.wait_for_termination()
+
+
 if __name__ == "__main__":
  if len(sys.argv) > 1:
  port = sys.argv[1]
  else:
  port = "50051"
  logging.basicConfig()
- serve(port)
+ if os.environ.get("GRPC_SERVER_ASYNC") == "1":
+ asyncio.run(serve_async(port))
+ else:
+ serve(port)

tests/contrib/grpc/grpc_client_tests.py

@@ -49,11 +49,14 @@
 from .grpc_app.testgrpc_pb2_grpc import TestServiceStub
 
 
-@pytest.fixture()
-def grpc_server(validating_httpserver, request):
+def setup_env(request, validating_httpserver):
  config = getattr(request, "param", {})
  env = {f"ELASTIC_APM_{k.upper()}": str(v) for k, v in config.items()}
  env.setdefault("ELASTIC_APM_SERVER_URL", validating_httpserver.url)
+ return env
+
+
+def setup_grpc_server(env):
  free_port = get_free_port()
  server_proc = subprocess.Popen(
  [os.path.join(sys.prefix, "bin", "python"), "-m", "tests.contrib.grpc.grpc_app.server", str(free_port)],
@@ -62,15 +65,32 @@ def grpc_server(validating_httpserver, request):
  env=env,
  )
  wait_for_open_port(free_port)
- yield f"localhost:{free_port}"
- server_proc.terminate()
+ return server_proc, free_port
 
 
 @pytest.fixture()
-def grpc_client_and_server_url(grpc_server):
- test_channel = grpc.insecure_channel(grpc_server)
+def env_fixture(validating_httpserver, request):
+ env = setup_env(request, validating_httpserver)
+ return env
+
+
+if hasattr(grpc, "aio"):
+ grpc_server_fixture_params = ["async", "sync"]
+else:
+ grpc_server_fixture_params = ["sync"]
+
+
+@pytest.fixture(params=grpc_server_fixture_params)
+def grpc_client_and_server_url(env_fixture, request):
+ env = {k: v for k, v in env_fixture.items()}
+ if request.param == "async":
+ env["GRPC_SERVER_ASYNC"] = "1"
+ server_proc, free_port = setup_grpc_server(env)
+ server_addr = f"localhost:{free_port}"
+ test_channel = grpc.insecure_channel(server_addr)
  test_client = TestServiceStub(test_channel)
- yield test_client, grpc_server
+ yield test_client, server_addr
+ server_proc.terminate()
 
 
 def test_grpc_client_server_instrumentation(instrument, sending_elasticapm_client, grpc_client_and_server_url):
@@ -200,7 +220,7 @@ def test_grpc_client_unsampled_span(instrument, sending_elasticapm_client, grpc_
 
 
 @pytest.mark.parametrize(
- "grpc_server",
+ "env_fixture",
  [
  {
  "recording": "False",