From 5763eb3dd46bee45c98cbaf1e118a6b01bbbc58d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Robert=20Paj=C4=85k?= <pellared@hotmail.com>
Date: Fri, 10 Nov 2023 12:39:27 +0100
Subject: [PATCH] Add BenchmarkStreamClientInterceptor

---
 .../otelgrpc/test/grpc_stats_handler_test.go  |  5 ++--
 .../grpc/otelgrpc/test/grpc_test.go           | 27 +++++++++----------
 .../grpc/otelgrpc/test/interceptor_test.go    | 18 +++++++++++++
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_stats_handler_test.go b/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_stats_handler_test.go
index 6768dfcb4d6..e6fd212f904 100644
--- a/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_stats_handler_test.go
+++ b/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_stats_handler_test.go
@@ -49,8 +49,7 @@ func TestStatsHandler(t *testing.T) {
 
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err, "failed to open port")
-	err = newGrpcTest(
-		listener,
+	client := newGrpcTest(t, listener,
 		[]grpc.DialOption{
 			grpc.WithStatsHandler(otelgrpc.NewClientHandler(
 				otelgrpc.WithTracerProvider(clientTP),
@@ -66,7 +65,7 @@ func TestStatsHandler(t *testing.T) {
 			),
 		},
 	)
-	require.NoError(t, err)
+	doCalls(client)
 
 	t.Run("ClientSpans", func(t *testing.T) {
 		checkClientSpans(t, clientSR.Ended())
diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_test.go b/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_test.go
index 9c99abf4b02..faec8107016 100644
--- a/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_test.go
+++ b/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_test.go
@@ -45,14 +45,18 @@ var wantInstrumentationScope = instrumentation.Scope{
 	Version:   otelgrpc.Version(),
 }
 
-// newGrpcTest creats a grpc server, starts it, and executes all the calls, closes everything down.
-func newGrpcTest(listener net.Listener, cOpt []grpc.DialOption, sOpt []grpc.ServerOption) error {
+// newGrpcTest creats a grpc server, starts it, and returns the client, closes everything down during test cleanup.
+func newGrpcTest(t testing.TB, listener net.Listener, cOpt []grpc.DialOption, sOpt []grpc.ServerOption) pb.TestServiceClient {
 	grpcServer := grpc.NewServer(sOpt...)
 	pb.RegisterTestServiceServer(grpcServer, interop.NewTestServer())
 	errCh := make(chan error)
 	go func() {
 		errCh <- grpcServer.Serve(listener)
 	}()
+	t.Cleanup(func() {
+		grpcServer.Stop()
+		assert.NoError(t, <-errCh)
+	})
 	ctx := context.Background()
 
 	cOpt = append(cOpt, grpc.WithBlock(), grpc.WithTransportCredentials(insecure.NewCredentials()))
@@ -67,17 +71,12 @@ func newGrpcTest(listener net.Listener, cOpt []grpc.DialOption, sOpt []grpc.Serv
 		listener.Addr().String(),
 		cOpt...,
 	)
-	if err != nil {
-		return err
-	}
-	client := pb.NewTestServiceClient(conn)
-
-	doCalls(client)
-
-	conn.Close()
-	grpcServer.Stop()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		assert.NoError(t, conn.Close())
+	})
 
-	return <-errCh
+	return pb.NewTestServiceClient(conn)
 }
 
 func doCalls(client pb.TestServiceClient) {
@@ -105,7 +104,7 @@ func TestInterceptors(t *testing.T) {
 
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err, "failed to open port")
-	err = newGrpcTest(listener,
+	client := newGrpcTest(t, listener,
 		[]grpc.DialOption{
 			//nolint:staticcheck // Interceptors are deprecated and will be removed in the next release.
 			grpc.WithUnaryInterceptor(otelgrpc.UnaryClientInterceptor(
@@ -132,7 +131,7 @@ func TestInterceptors(t *testing.T) {
 			)),
 		},
 	)
-	require.NoError(t, err)
+	doCalls(client)
 
 	t.Run("UnaryClientSpans", func(t *testing.T) {
 		checkUnaryClientSpans(t, clientUnarySR.Ended(), listener.Addr().String())
diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/test/interceptor_test.go b/instrumentation/google.golang.org/grpc/otelgrpc/test/interceptor_test.go
index 400f970dba0..e4714211caa 100644
--- a/instrumentation/google.golang.org/grpc/otelgrpc/test/interceptor_test.go
+++ b/instrumentation/google.golang.org/grpc/otelgrpc/test/interceptor_test.go
@@ -40,6 +40,7 @@ import (
 	"google.golang.org/grpc"
 	grpc_codes "google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/interop"
 	"google.golang.org/grpc/interop/grpc_testing"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
@@ -1128,3 +1129,20 @@ func assertServerMetrics(t *testing.T, reader metric.Reader, serviceName, name s
 	require.Len(t, rm.ScopeMetrics, 1)
 	metricdatatest.AssertEqual(t, want, rm.ScopeMetrics[0], metricdatatest.IgnoreTimestamp(), metricdatatest.IgnoreValue())
 }
+
+func BenchmarkStreamClientInterceptor(b *testing.B) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(b, err, "failed to open port")
+	client := newGrpcTest(b, listener,
+		[]grpc.DialOption{
+			//nolint:staticcheck // Interceptors are deprecated and will be removed in the next release.
+			grpc.WithStreamInterceptor(otelgrpc.StreamClientInterceptor()),
+		},
+		[]grpc.ServerOption{},
+	)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		interop.DoClientStreaming(client)
+	}
+}