From 53370433e2de78cc1bc9f506b42735282535f7f8 Mon Sep 17 00:00:00 2001
From: tmdiep <tmdiep@google.com>
Date: Sat, 5 Dec 2020 07:10:03 +1100
Subject: [PATCH] test(pubsublite): refactor out common code for main_test.go
 (#3276)

Since main_test.go needs to be cloned into every package that needs to use the mock server, refactor the bulk of the common code to keep it minimal.
---
 pubsublite/internal/test/mock.go      | 33 ++++++++++++++++++---------
 pubsublite/internal/wire/main_test.go | 13 ++---------
 pubsublite/main_test.go               | 13 ++---------
 3 files changed, 26 insertions(+), 33 deletions(-)

diff --git a/pubsublite/internal/test/mock.go b/pubsublite/internal/test/mock.go
index 1ddfa91e538..515fbec2668 100644
--- a/pubsublite/internal/test/mock.go
+++ b/pubsublite/internal/test/mock.go
@@ -16,11 +16,13 @@ package test
 import (
 	"context"
 	"io"
+	"log"
 	"reflect"
 	"sync"
 
 	"cloud.google.com/go/internal/testutil"
 	"cloud.google.com/go/internal/uid"
+	"google.golang.org/api/option"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -29,12 +31,6 @@ import (
 	pb "google.golang.org/genproto/googleapis/cloud/pubsublite/v1"
 )
 
-// Server is a mock Pub/Sub Lite server that can be used for unit testing.
-type Server struct {
-	LiteServer MockServer
-	gRPCServer *testutil.Server
-}
-
 // MockServer is an in-memory mock implementation of a Pub/Sub Lite service,
 // which allows unit tests to inspect requests received by the server and send
 // fake responses.
@@ -48,6 +44,12 @@ type MockServer interface {
 	OnTestEnd()
 }
 
+// Server is a mock Pub/Sub Lite server that can be used for unit testing.
+type Server struct {
+	LiteServer MockServer
+	gRPCServer *testutil.Server
+}
+
 // NewServer creates a new mock Pub/Sub Lite server.
 func NewServer() (*Server, error) {
 	srv, err := testutil.NewServer()
@@ -64,6 +66,20 @@ func NewServer() (*Server, error) {
 	return &Server{LiteServer: liteServer, gRPCServer: srv}, nil
 }
 
+// NewServerWithConn creates a new mock Pub/Sub Lite server along with client
+// options to connect to it.
+func NewServerWithConn() (*Server, []option.ClientOption) {
+	testServer, err := NewServer()
+	if err != nil {
+		log.Fatal(err)
+	}
+	conn, err := grpc.Dial(testServer.Addr(), grpc.WithInsecure())
+	if err != nil {
+		log.Fatal(err)
+	}
+	return testServer, []option.ClientOption{option.WithGRPCConn(conn)}
+}
+
 // Addr returns the address that the server is listening on.
 func (s *Server) Addr() string {
 	return s.gRPCServer.Addr
@@ -74,11 +90,6 @@ func (s *Server) Close() {
 	s.gRPCServer.Close()
 }
 
-type streamHolder struct {
-	stream   grpc.ServerStream
-	verifier *RPCVerifier
-}
-
 // mockLiteServer implements the MockServer interface.
 type mockLiteServer struct {
 	pb.AdminServiceServer
diff --git a/pubsublite/internal/wire/main_test.go b/pubsublite/internal/wire/main_test.go
index 0f9b401c477..a9dbc679194 100644
--- a/pubsublite/internal/wire/main_test.go
+++ b/pubsublite/internal/wire/main_test.go
@@ -15,13 +15,11 @@ package wire
 
 import (
 	"flag"
-	"log"
 	"os"
 	"testing"
 
 	"cloud.google.com/go/pubsublite/internal/test"
 	"google.golang.org/api/option"
-	"google.golang.org/grpc"
 )
 
 var (
@@ -33,16 +31,9 @@ var (
 func TestMain(m *testing.M) {
 	flag.Parse()
 
-	testServer, err := test.NewServer()
-	if err != nil {
-		log.Fatal(err)
-	}
+	testServer, clientOpts := test.NewServerWithConn()
 	mockServer = testServer.LiteServer
-	conn, err := grpc.Dial(testServer.Addr(), grpc.WithInsecure())
-	if err != nil {
-		log.Fatal(err)
-	}
-	testClientOpts = []option.ClientOption{option.WithGRPCConn(conn)}
+	testClientOpts = clientOpts
 
 	exit := m.Run()
 	testServer.Close()
diff --git a/pubsublite/main_test.go b/pubsublite/main_test.go
index e77724676e8..0f325f8d904 100644
--- a/pubsublite/main_test.go
+++ b/pubsublite/main_test.go
@@ -15,13 +15,11 @@ package pubsublite
 
 import (
 	"flag"
-	"log"
 	"os"
 	"testing"
 
 	"cloud.google.com/go/pubsublite/internal/test"
 	"google.golang.org/api/option"
-	"google.golang.org/grpc"
 )
 
 var (
@@ -33,16 +31,9 @@ var (
 func TestMain(m *testing.M) {
 	flag.Parse()
 
-	testServer, err := test.NewServer()
-	if err != nil {
-		log.Fatal(err)
-	}
+	testServer, clientOpts := test.NewServerWithConn()
 	mockServer = testServer.LiteServer
-	conn, err := grpc.Dial(testServer.Addr(), grpc.WithInsecure())
-	if err != nil {
-		log.Fatal(err)
-	}
-	testClientOpts = []option.ClientOption{option.WithGRPCConn(conn)}
+	testClientOpts = clientOpts
 
 	exit := m.Run()
 	testServer.Close()