test: switching to stubserver in tests instead of testservice (#7925)

grpc · Dec 18, 2024 · b3bdacb · b3bdacb
1 parent e8055ea
commit b3bdacb
Showing 1 changed file with 31 additions and 21 deletions.
diff --git a/test/balancer_test.go b/test/balancer_test.go
@@ -417,11 +417,15 @@ func (s) TestAddressAttributesInNewSubConn(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-
-	s := grpc.NewServer()
-	testgrpc.RegisterTestServiceServer(s, &testServer{})
-	go s.Serve(lis)
-	defer s.Stop()
+	stub := &stubserver.StubServer{
+		Listener: lis,
+		EmptyCallF: func(_ context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
+			return &testpb.Empty{}, nil
+		},
+		S: grpc.NewServer(),
+	}
+	stubserver.StartTestService(t, stub)
+	defer stub.S.Stop()
 	t.Logf("Started gRPC server at %s...", lis.Addr().String())
 
 	creds := &attrTransportCreds{}
@@ -548,15 +552,16 @@ func (s) TestServersSwap(t *testing.T) {
 		if err != nil {
 			t.Fatalf("Error while listening. Err: %v", err)
 		}
-		s := grpc.NewServer()
-		ts := &funcServer{
-			unaryCall: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+
+		stub := &stubserver.StubServer{
+			Listener: lis,
+			UnaryCallF: func(_ context.Context, _ *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
 				return &testpb.SimpleResponse{Username: username}, nil
 			},
+			S: grpc.NewServer(),
 		}
-		testgrpc.RegisterTestServiceServer(s, ts)
-		go s.Serve(lis)
-		return lis.Addr().String(), s.Stop
+		stubserver.StartTestService(t, stub)
+		return lis.Addr().String(), stub.S.Stop
 	}
 	const one = "1"
 	addr1, cleanup := reg(one)
@@ -603,16 +608,16 @@ func (s) TestWaitForReady(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Error while listening. Err: %v", err)
 	}
-	s := grpc.NewServer()
-	defer s.Stop()
 	const one = "1"
-	ts := &funcServer{
-		unaryCall: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+	stub := &stubserver.StubServer{
+		Listener: lis,
+		UnaryCallF: func(_ context.Context, _ *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
 			return &testpb.SimpleResponse{Username: one}, nil
 		},
+		S: grpc.NewServer(),
 	}
-	testgrpc.RegisterTestServiceServer(s, ts)
-	go s.Serve(lis)
+	stubserver.StartTestService(t, stub)
+	defer stub.S.Stop()
 
 	// Initialize client
 	r := manual.NewBuilderWithScheme("whatever")
@@ -740,10 +745,15 @@ func (s) TestAuthorityInBuildOptions(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			s := grpc.NewServer()
-			testgrpc.RegisterTestServiceServer(s, &testServer{})
-			go s.Serve(lis)
-			defer s.Stop()
+			stub := &stubserver.StubServer{
+				Listener: lis,
+				EmptyCallF: func(_ context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
+					return &testpb.Empty{}, nil
+				},
+				S: grpc.NewServer(),
+			}
+			stubserver.StartTestService(t, stub)
+			defer stub.S.Stop()
 			t.Logf("Started gRPC server at %s...", lis.Addr().String())
 
 			r := manual.NewBuilderWithScheme("whatever")