From 00145448aa742a1ae0a0b14aff1d08100dfc676e Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Thu, 18 Mar 2021 15:29:42 -0700
Subject: [PATCH] [admin_api] t refactoring

---
 admin/admin_test.go      | 20 +++++---------------
 admin/test/admin_test.go | 18 ++++--------------
 admin/test/utils.go      | 37 +++++++++++++++++--------------------
 3 files changed, 26 insertions(+), 49 deletions(-)

diff --git a/admin/admin_test.go b/admin/admin_test.go
index 470331635da0..0ee4aade0f3c 100644
--- a/admin/admin_test.go
+++ b/admin/admin_test.go
@@ -26,19 +26,9 @@ import (
 )
 
 func TestRegisterNoCSDS(t *testing.T) {
-	if err := test.RunRegisterTests([]test.RegisterTestCase{
-		{
-			Name: "channelz",
-			Run:  test.RunChannelz,
-			Code: codes.OK,
-		},
-		{
-			Name: "CSDS",
-			Run:  test.RunCSDS,
-			// CSDS is not registered because xDS isn't imported.
-			Code: codes.Unimplemented,
-		},
-	}); err != nil {
-		t.Fatal(err)
-	}
+	test.RunRegisterTests(t, test.ExpectedStatusCodes{
+		ChannelzCode: codes.OK,
+		// CSDS is not registered because xDS isn't imported.
+		CSDSCode: codes.Unimplemented,
+	})
 }
diff --git a/admin/test/admin_test.go b/admin/test/admin_test.go
index 232d915c6467..f0f784bfdf36 100644
--- a/admin/test/admin_test.go
+++ b/admin/test/admin_test.go
@@ -31,18 +31,8 @@ import (
 )
 
 func TestRegisterWithCSDS(t *testing.T) {
-	if err := test.RunRegisterTests([]test.RegisterTestCase{
-		{
-			Name: "channelz",
-			Run:  test.RunChannelz,
-			Code: codes.OK,
-		},
-		{
-			Name: "CSDS",
-			Run:  test.RunCSDS,
-			Code: codes.OK,
-		},
-	}); err != nil {
-		t.Fatal(err)
-	}
+	test.RunRegisterTests(t, test.ExpectedStatusCodes{
+		ChannelzCode: codes.OK,
+		CSDSCode:     codes.OK,
+	})
 }
diff --git a/admin/test/utils.go b/admin/test/utils.go
index b97b61bedef4..9187ee659623 100644
--- a/admin/test/utils.go
+++ b/admin/test/utils.go
@@ -22,8 +22,8 @@ package test
 
 import (
 	"context"
-	"fmt"
 	"net"
+	"testing"
 	"time"
 
 	v3statuspb "github.com/envoyproxy/go-control-plane/envoy/service/status/v3"
@@ -40,20 +40,16 @@ const (
 	defaultTestTimeout = 10 * time.Second
 )
 
-// RegisterTestCase contains the function to make an RPC and the expected status code
-// (can be OK).
-type RegisterTestCase struct {
-	// Name of this RPC.
-	Name string
-	// The function to make the RPC on the ClientConn.
-	Run func(*grpc.ClientConn) error
-	// The expected code returned from the RPC.
-	Code codes.Code
+// ExpectedStatusCodes contains the expected status code for each RPC (can be
+// OK).
+type ExpectedStatusCodes struct {
+	ChannelzCode codes.Code
+	CSDSCode     codes.Code
 }
 
 // RunRegisterTests makes a client, runs the RPCs, and compares the status
 // codes.
-func RunRegisterTests(rcs []RegisterTestCase) error {
+func RunRegisterTests(t *testing.T, ec ExpectedStatusCodes) {
 	nodeID := uuid.New().String()
 	bootstrapCleanup, err := xds.SetupBootstrapFile(xds.BootstrapOptions{
 		Version:   xds.TransportV3,
@@ -61,20 +57,20 @@ func RunRegisterTests(rcs []RegisterTestCase) error {
 		ServerURI: "no.need.for.a.server",
 	})
 	if err != nil {
-		return err
+		t.Fatal(err)
 	}
 	defer bootstrapCleanup()
 
 	lis, err := net.Listen("tcp", "localhost:0")
 	if err != nil {
-		return fmt.Errorf("cannot create listener: %v", err)
+		t.Fatalf("cannot create listener: %v", err)
 	}
 
 	server := grpc.NewServer()
 	defer server.Stop()
 	cleanup, err := admin.Register(server)
 	if err != nil {
-		return fmt.Errorf("failed to register admin: %v", err)
+		t.Fatalf("failed to register admin: %v", err)
 	}
 	defer cleanup()
 	go func() {
@@ -83,15 +79,16 @@ func RunRegisterTests(rcs []RegisterTestCase) error {
 
 	conn, err := grpc.Dial(lis.Addr().String(), grpc.WithInsecure())
 	if err != nil {
-		return fmt.Errorf("cannot connect to server: %v", err)
+		t.Fatalf("cannot connect to server: %v", err)
 	}
 
-	for _, rc := range rcs {
-		if err := rc.Run(conn); status.Code(err) != rc.Code {
-			return fmt.Errorf("%s test failed with error %v, want code %v", rc.Name, err, rc.Code)
-		}
+	if err := RunChannelz(conn); status.Code(err) != ec.ChannelzCode {
+		t.Fatalf("%s test failed with error %v, want code %v", "channelz", err, ec.ChannelzCode)
+	}
+
+	if err := RunCSDS(conn); status.Code(err) != ec.CSDSCode {
+		t.Fatalf("%s test failed with error %v, want code %v", "CSDS", err, ec.CSDSCode)
 	}
-	return nil
 }
 
 // RunChannelz makes a channelz RPC.