From 2fe71180762478c66b0027f780b95c40fc563a55 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Mon, 11 Oct 2021 15:42:10 -0700
Subject: [PATCH] xds/e2e: move flag check to each test, and call t.Skip()
 (#4861)

---
 xds/internal/test/e2e/e2e.go      |  8 +++++---
 xds/internal/test/e2e/e2e_test.go | 17 ++++++-----------
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/xds/internal/test/e2e/e2e.go b/xds/internal/test/e2e/e2e.go
index 82c7f9dfb..ade6339bf 100644
--- a/xds/internal/test/e2e/e2e.go
+++ b/xds/internal/test/e2e/e2e.go
@@ -26,7 +26,9 @@ import (
 	"os/exec"
 
 	"google.golang.org/grpc"
+	channelzgrpc "google.golang.org/grpc/channelz/grpc_channelz_v1"
 	channelzpb "google.golang.org/grpc/channelz/grpc_channelz_v1"
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
 	testpb "google.golang.org/grpc/interop/grpc_testing"
 )
 
@@ -83,7 +85,7 @@ func newClient(target, binaryPath, bootstrap string, logger io.Writer, flags ...
 }
 
 func (c *client) clientStats(ctx context.Context) (*testpb.LoadBalancerStatsResponse, error) {
-	ccc := testpb.NewLoadBalancerStatsServiceClient(c.statsCC)
+	ccc := testgrpc.NewLoadBalancerStatsServiceClient(c.statsCC)
 	return ccc.GetClientStats(ctx, &testpb.LoadBalancerStatsRequest{
 		NumRpcs:    100,
 		TimeoutSec: 10,
@@ -91,13 +93,13 @@ func (c *client) clientStats(ctx context.Context) (*testpb.LoadBalancerStatsResp
 }
 
 func (c *client) configRPCs(ctx context.Context, req *testpb.ClientConfigureRequest) error {
-	ccc := testpb.NewXdsUpdateClientConfigureServiceClient(c.statsCC)
+	ccc := testgrpc.NewXdsUpdateClientConfigureServiceClient(c.statsCC)
 	_, err := ccc.Configure(ctx, req)
 	return err
 }
 
 func (c *client) channelzSubChannels(ctx context.Context) ([]*channelzpb.Subchannel, error) {
-	ccc := channelzpb.NewChannelzClient(c.statsCC)
+	ccc := channelzgrpc.NewChannelzClient(c.statsCC)
 	r, err := ccc.GetTopChannels(ctx, &channelzpb.GetTopChannelsRequest{})
 	if err != nil {
 		return nil, err
diff --git a/xds/internal/test/e2e/e2e_test.go b/xds/internal/test/e2e/e2e_test.go
index 5c116b478..6984566db 100644
--- a/xds/internal/test/e2e/e2e_test.go
+++ b/xds/internal/test/e2e/e2e_test.go
@@ -39,17 +39,6 @@ var (
 	serverPath = flag.String("server", "./binaries/server", "The interop server")
 )
 
-func TestMain(m *testing.M) {
-	flag.Parse()
-	if _, err := os.Stat(*clientPath); os.IsNotExist(err) {
-		return
-	}
-	if _, err := os.Stat(*serverPath); os.IsNotExist(err) {
-		return
-	}
-	os.Exit(m.Run())
-}
-
 type testOpts struct {
 	testName     string
 	backendCount int
@@ -58,6 +47,12 @@ type testOpts struct {
 
 func setup(t *testing.T, opts testOpts) (*controlPlane, *client, []*server) {
 	t.Helper()
+	if _, err := os.Stat(*clientPath); os.IsNotExist(err) {
+		t.Skip("skipped because client is not found")
+	}
+	if _, err := os.Stat(*serverPath); os.IsNotExist(err) {
+		t.Skip("skipped because server is not found")
+	}
 	backendCount := 1
 	if opts.backendCount != 0 {
 		backendCount = opts.backendCount