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

@@ -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

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()

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()