Repository: grpc/grpc-go
File: balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go (Go)

Code before:
func Test(t *testing.T) {
	grpctest.RunSubTests(t, s{})
}

// testServer is a server than can be stopped and resumed without closing
// the listener. This guarantees the same port number (and address) is used
// after restart. When a server is stopped, it accepts and closes all tcp
// connections from clients.
type testServer struct {
	stubserver.StubServer
	lis *testutils.RestartableListener
}

func (s *testServer) stop() {
	s.lis.Stop()
}

func (s *testServer) resume() {
	s.lis.Restart()
}

func newTestServer() (*testServer, error) {
	l, err := testutils.LocalTCPListener()
	if err != nil {
		return nil, err
	}
	rl := testutils.NewRestartableListener(l)
	ss := stubserver.StubServer{
		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) { return &testpb.Empty{}, nil },
		Listener:   rl,
	}
	return &testServer{
		StubServer: ss,
		lis:        rl,
	}, nil
}

// setupPickFirstLeaf performs steps required for pick_first tests. It starts a
// bunch of backends exporting the TestService, and creates a ClientConn to them.
func setupPickFirstLeaf(t *testing.T, backendCount int, opts ...grpc.DialOption) (*grpc.ClientConn, *manual.Resolver, *backendManager) {
	t.Helper()
	r := manual.NewBuilderWithScheme("whatever")
	backends := make([]*testServer, backendCount)
	addrs := make([]resolver.Address, backendCount)

	for i := 0; i < backendCount; i++ {
		server, err := newTestServer()
		if err != nil {
			t.Fatalf("net.Listen() failed: %v", err)
		}
		backend := stubserver.StartTestService(t, &server.StubServer)

Review comment (bug):
We *could* accept a `testing.T` here, make this a test helper and call `t.Fatalf` here if creating the listener fails. That would simplify callsites since they wont have to handle two return values.

Code after:
func Test(t *testing.T) {
	grpctest.RunSubTests(t, s{})
}

// testServer is a server than can be stopped and resumed without closing
// the listener. This guarantees the same port number (and address) is used
// after restart. When a server is stopped, it accepts and closes all tcp
// connections from clients.
type testServer struct {
	stubserver.StubServer
	lis *testutils.RestartableListener
}

func (s *testServer) stop() {
	s.lis.Stop()
}

func (s *testServer) resume() {
	s.lis.Restart()
}

func newTestServer(t *testing.T) *testServer {
	l, err := testutils.LocalTCPListener()
	if err != nil {
		t.Fatalf("Failed to create listener: %v", err)
	}
	rl := testutils.NewRestartableListener(l)
	ss := stubserver.StubServer{
		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) { return &testpb.Empty{}, nil },
		Listener:   rl,
	}
	return &testServer{
		StubServer: ss,
		lis:        rl,
	}
}

// setupPickFirstLeaf performs steps required for pick_first tests. It starts a
// bunch of backends exporting the TestService, and creates a ClientConn to them.
func setupPickFirstLeaf(t *testing.T, backendCount int, opts ...grpc.DialOption) (*grpc.ClientConn, *manual.Resolver, *backendManager) {
	t.Helper()
	r := manual.NewBuilderWithScheme("whatever")
	backends := make([]*testServer, backendCount)
	addrs := make([]resolver.Address, backendCount)

	for i := 0; i < backendCount; i++ {
		server := newTestServer(t)
		backend := stubserver.StartTestService(t, &server.StubServer)
		t.Cleanup(func() {
			backend.Stop()
		})