service: remove stop method and use contexts (#7292)

privval/grpc/util.go

@@ -109,7 +109,7 @@ func DialRemoteSigner(
 
 	dialOptions = append(dialOptions, transportSecurity)
 
-	ctx := context.Background()
+	ctx := context.TODO()
 	_, address := tmnet.ProtocolAndAddress(cfg.ListenAddr)
 	conn, err := grpc.DialContext(ctx, address, dialOptions...)
 	if err != nil {

privval/signer_client.go

@@ -23,9 +23,9 @@ var _ types.PrivValidator = (*SignerClient)(nil)
 
 // NewSignerClient returns an instance of SignerClient.
 // it will start the endpoint (if not already started)
-func NewSignerClient(endpoint *SignerListenerEndpoint, chainID string) (*SignerClient, error) {
+func NewSignerClient(ctx context.Context, endpoint *SignerListenerEndpoint, chainID string) (*SignerClient, error) {
 	if !endpoint.IsRunning() {
-		if err := endpoint.Start(); err != nil {
+		if err := endpoint.Start(ctx); err != nil {
 			return nil, fmt.Errorf("failed to start listener endpoint: %w", err)
 		}
 	}

privval/signer_client_test.go

@@ -23,370 +23,336 @@ type signerTestCase struct {
 	mockPV       types.PrivValidator
 	signerClient *SignerClient
 	signerServer *SignerServer
+	name         string
+	closer       context.CancelFunc
 }
 
-func getSignerTestCases(t *testing.T) []signerTestCase {
+func getSignerTestCases(ctx context.Context, t *testing.T) []signerTestCase {
+	t.Helper()
+
 	testCases := make([]signerTestCase, 0)
 
 	// Get test cases for each possible dialer (DialTCP / DialUnix / etc)
-	for _, dtc := range getDialerTestCases(t) {
+	for idx, dtc := range getDialerTestCases(t) {
 		chainID := tmrand.Str(12)
 		mockPV := types.NewMockPV()
 
+		cctx, ccancel := context.WithCancel(ctx)
 		// get a pair of signer listener, signer dialer endpoints
-		sl, sd := getMockEndpoints(t, dtc.addr, dtc.dialer)
-		sc, err := NewSignerClient(sl, chainID)
+		sl, sd := getMockEndpoints(cctx, t, dtc.addr, dtc.dialer)
+		sc, err := NewSignerClient(cctx, sl, chainID)
 		require.NoError(t, err)
 		ss := NewSignerServer(sd, chainID, mockPV)
 
-		err = ss.Start()
-		require.NoError(t, err)
+		require.NoError(t, ss.Start(cctx))
 
-		tc := signerTestCase{
+		testCases = append(testCases, signerTestCase{
+			name:         fmt.Sprintf("Case%d%T_%s", idx, dtc.dialer, chainID),
+			closer:       ccancel,
 			chainID:      chainID,
 			mockPV:       mockPV,
 			signerClient: sc,
 			signerServer: ss,
-		}
-
-		testCases = append(testCases, tc)
+		})
 	}
 
 	return testCases
 }
 
 func TestSignerClose(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		err := tc.signerClient.Close()
-		assert.NoError(t, err)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		err = tc.signerServer.Stop()
-		assert.NoError(t, err)
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
+
+			assert.NoError(t, tc.signerClient.Close())
+			assert.NoError(t, tc.signerServer.Stop())
+		})
 	}
 }
 
 func TestSignerPing(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
-			}
-		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
-			}
-		})
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
+	for _, tc := range getSignerTestCases(ctx, t) {
 		err := tc.signerClient.Ping()
 		assert.NoError(t, err)
 	}
 }
 
 func TestSignerGetPubKey(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
-			}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
+
+			pubKey, err := tc.signerClient.GetPubKey(ctx)
+			require.NoError(t, err)
+			expectedPubKey, err := tc.mockPV.GetPubKey(ctx)
+			require.NoError(t, err)
+
+			assert.Equal(t, expectedPubKey, pubKey)
+
+			pubKey, err = tc.signerClient.GetPubKey(ctx)
+			require.NoError(t, err)
+			expectedpk, err := tc.mockPV.GetPubKey(ctx)
+			require.NoError(t, err)
+			expectedAddr := expectedpk.Address()
+
+			assert.Equal(t, expectedAddr, pubKey.Address())
 		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
-			}
-		})
-
-		pubKey, err := tc.signerClient.GetPubKey(context.Background())
-		require.NoError(t, err)
-		expectedPubKey, err := tc.mockPV.GetPubKey(context.Background())
-		require.NoError(t, err)
-
-		assert.Equal(t, expectedPubKey, pubKey)
-
-		pubKey, err = tc.signerClient.GetPubKey(context.Background())
-		require.NoError(t, err)
-		expectedpk, err := tc.mockPV.GetPubKey(context.Background())
-		require.NoError(t, err)
-		expectedAddr := expectedpk.Address()
-
-		assert.Equal(t, expectedAddr, pubKey.Address())
 	}
 }
 
 func TestSignerProposal(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		ts := time.Now()
-		hash := tmrand.Bytes(tmhash.Size)
-		have := &types.Proposal{
-			Type:      tmproto.ProposalType,
-			Height:    1,
-			Round:     2,
-			POLRound:  2,
-			BlockID:   types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp: ts,
-		}
-		want := &types.Proposal{
-			Type:      tmproto.ProposalType,
-			Height:    1,
-			Round:     2,
-			POLRound:  2,
-			BlockID:   types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp: ts,
-		}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
+
+			ts := time.Now()
+			hash := tmrand.Bytes(tmhash.Size)
+			have := &types.Proposal{
+				Type:      tmproto.ProposalType,
+				Height:    1,
+				Round:     2,
+				POLRound:  2,
+				BlockID:   types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp: ts,
 			}
-		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
+			want := &types.Proposal{
+				Type:      tmproto.ProposalType,
+				Height:    1,
+				Round:     2,
+				POLRound:  2,
+				BlockID:   types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp: ts,
 			}
+
+			require.NoError(t, tc.mockPV.SignProposal(ctx, tc.chainID, want.ToProto()))
+			require.NoError(t, tc.signerClient.SignProposal(ctx, tc.chainID, have.ToProto()))
+
+			assert.Equal(t, want.Signature, have.Signature)
 		})
 
-		require.NoError(t, tc.mockPV.SignProposal(context.Background(), tc.chainID, want.ToProto()))
-		require.NoError(t, tc.signerClient.SignProposal(context.Background(), tc.chainID, have.ToProto()))
-
-		assert.Equal(t, want.Signature, have.Signature)
 	}
 }
 
 func TestSignerVote(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		ts := time.Now()
-		hash := tmrand.Bytes(tmhash.Size)
-		valAddr := tmrand.Bytes(crypto.AddressSize)
-		want := &types.Vote{
-			Type:             tmproto.PrecommitType,
-			Height:           1,
-			Round:            2,
-			BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp:        ts,
-			ValidatorAddress: valAddr,
-			ValidatorIndex:   1,
-		}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		have := &types.Vote{
-			Type:             tmproto.PrecommitType,
-			Height:           1,
-			Round:            2,
-			BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp:        ts,
-			ValidatorAddress: valAddr,
-			ValidatorIndex:   1,
-		}
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
 
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
+			ts := time.Now()
+			hash := tmrand.Bytes(tmhash.Size)
+			valAddr := tmrand.Bytes(crypto.AddressSize)
+			want := &types.Vote{
+				Type:             tmproto.PrecommitType,
+				Height:           1,
+				Round:            2,
+				BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp:        ts,
+				ValidatorAddress: valAddr,
+				ValidatorIndex:   1,
 			}
-		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
+
+			have := &types.Vote{
+				Type:             tmproto.PrecommitType,
+				Height:           1,
+				Round:            2,
+				BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp:        ts,
+				ValidatorAddress: valAddr,
+				ValidatorIndex:   1,
 			}
+
+			require.NoError(t, tc.mockPV.SignVote(ctx, tc.chainID, want.ToProto()))
+			require.NoError(t, tc.signerClient.SignVote(ctx, tc.chainID, have.ToProto()))
+
+			assert.Equal(t, want.Signature, have.Signature)
 		})
-
-		require.NoError(t, tc.mockPV.SignVote(context.Background(), tc.chainID, want.ToProto()))
-		require.NoError(t, tc.signerClient.SignVote(context.Background(), tc.chainID, have.ToProto()))
-
-		assert.Equal(t, want.Signature, have.Signature)
 	}
 }
 
 func TestSignerVoteResetDeadline(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		ts := time.Now()
-		hash := tmrand.Bytes(tmhash.Size)
-		valAddr := tmrand.Bytes(crypto.AddressSize)
-		want := &types.Vote{
-			Type:             tmproto.PrecommitType,
-			Height:           1,
-			Round:            2,
-			BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp:        ts,
-			ValidatorAddress: valAddr,
-			ValidatorIndex:   1,
-		}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		have := &types.Vote{
-			Type:             tmproto.PrecommitType,
-			Height:           1,
-			Round:            2,
-			BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp:        ts,
-			ValidatorAddress: valAddr,
-			ValidatorIndex:   1,
-		}
-
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			ts := time.Now()
+			hash := tmrand.Bytes(tmhash.Size)
+			valAddr := tmrand.Bytes(crypto.AddressSize)
+			want := &types.Vote{
+				Type:             tmproto.PrecommitType,
+				Height:           1,
+				Round:            2,
+				BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp:        ts,
+				ValidatorAddress: valAddr,
+				ValidatorIndex:   1,
 			}
-		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
+
+			have := &types.Vote{
+				Type:             tmproto.PrecommitType,
+				Height:           1,
+				Round:            2,
+				BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp:        ts,
+				ValidatorAddress: valAddr,
+				ValidatorIndex:   1,
 			}
+
+			time.Sleep(testTimeoutReadWrite2o3)
+
+			require.NoError(t, tc.mockPV.SignVote(ctx, tc.chainID, want.ToProto()))
+			require.NoError(t, tc.signerClient.SignVote(ctx, tc.chainID, have.ToProto()))
+			assert.Equal(t, want.Signature, have.Signature)
+
+			// TODO(jleni): Clarify what is actually being tested
+
+			// This would exceed the deadline if it was not extended by the previous message
+			time.Sleep(testTimeoutReadWrite2o3)
+
+			require.NoError(t, tc.mockPV.SignVote(ctx, tc.chainID, want.ToProto()))
+			require.NoError(t, tc.signerClient.SignVote(ctx, tc.chainID, have.ToProto()))
+			assert.Equal(t, want.Signature, have.Signature)
 		})
-
-		time.Sleep(testTimeoutReadWrite2o3)
-
-		require.NoError(t, tc.mockPV.SignVote(context.Background(), tc.chainID, want.ToProto()))
-		require.NoError(t, tc.signerClient.SignVote(context.Background(), tc.chainID, have.ToProto()))
-		assert.Equal(t, want.Signature, have.Signature)
-
-		// TODO(jleni): Clarify what is actually being tested
-
-		// This would exceed the deadline if it was not extended by the previous message
-		time.Sleep(testTimeoutReadWrite2o3)
-
-		require.NoError(t, tc.mockPV.SignVote(context.Background(), tc.chainID, want.ToProto()))
-		require.NoError(t, tc.signerClient.SignVote(context.Background(), tc.chainID, have.ToProto()))
-		assert.Equal(t, want.Signature, have.Signature)
 	}
 }
 
 func TestSignerVoteKeepAlive(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		ts := time.Now()
-		hash := tmrand.Bytes(tmhash.Size)
-		valAddr := tmrand.Bytes(crypto.AddressSize)
-		want := &types.Vote{
-			Type:             tmproto.PrecommitType,
-			Height:           1,
-			Round:            2,
-			BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp:        ts,
-			ValidatorAddress: valAddr,
-			ValidatorIndex:   1,
-		}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		have := &types.Vote{
-			Type:             tmproto.PrecommitType,
-			Height:           1,
-			Round:            2,
-			BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp:        ts,
-			ValidatorAddress: valAddr,
-			ValidatorIndex:   1,
-		}
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
 
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
+			ts := time.Now()
+			hash := tmrand.Bytes(tmhash.Size)
+			valAddr := tmrand.Bytes(crypto.AddressSize)
+			want := &types.Vote{
+				Type:             tmproto.PrecommitType,
+				Height:           1,
+				Round:            2,
+				BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp:        ts,
+				ValidatorAddress: valAddr,
+				ValidatorIndex:   1,
 			}
-		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
+
+			have := &types.Vote{
+				Type:             tmproto.PrecommitType,
+				Height:           1,
+				Round:            2,
+				BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp:        ts,
+				ValidatorAddress: valAddr,
+				ValidatorIndex:   1,
 			}
+
+			// Check that even if the client does not request a
+			// signature for a long time. The service is still available
+
+			// in this particular case, we use the dialer logger to ensure that
+			// test messages are properly interleaved in the test logs
+			tc.signerServer.Logger.Debug("TEST: Forced Wait -------------------------------------------------")
+			time.Sleep(testTimeoutReadWrite * 3)
+			tc.signerServer.Logger.Debug("TEST: Forced Wait DONE---------------------------------------------")
+
+			require.NoError(t, tc.mockPV.SignVote(ctx, tc.chainID, want.ToProto()))
+			require.NoError(t, tc.signerClient.SignVote(ctx, tc.chainID, have.ToProto()))
+
+			assert.Equal(t, want.Signature, have.Signature)
 		})
-
-		// Check that even if the client does not request a
-		// signature for a long time. The service is still available
-
-		// in this particular case, we use the dialer logger to ensure that
-		// test messages are properly interleaved in the test logs
-		tc.signerServer.Logger.Debug("TEST: Forced Wait -------------------------------------------------")
-		time.Sleep(testTimeoutReadWrite * 3)
-		tc.signerServer.Logger.Debug("TEST: Forced Wait DONE---------------------------------------------")
-
-		require.NoError(t, tc.mockPV.SignVote(context.Background(), tc.chainID, want.ToProto()))
-		require.NoError(t, tc.signerClient.SignVote(context.Background(), tc.chainID, have.ToProto()))
-
-		assert.Equal(t, want.Signature, have.Signature)
 	}
 }
 
 func TestSignerSignProposalErrors(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		// Replace service with a mock that always fails
-		tc.signerServer.privVal = types.NewErroringMockPV()
-		tc.mockPV = types.NewErroringMockPV()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
+			// Replace service with a mock that always fails
+			tc.signerServer.privVal = types.NewErroringMockPV()
+			tc.mockPV = types.NewErroringMockPV()
+
+			ts := time.Now()
+			hash := tmrand.Bytes(tmhash.Size)
+			proposal := &types.Proposal{
+				Type:      tmproto.ProposalType,
+				Height:    1,
+				Round:     2,
+				POLRound:  2,
+				BlockID:   types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp: ts,
+				Signature: []byte("signature"),
 			}
+
+			err := tc.signerClient.SignProposal(ctx, tc.chainID, proposal.ToProto())
+			rserr, ok := err.(*RemoteSignerError)
+			require.True(t, ok, "%T", err)
+			require.Contains(t, rserr.Error(), types.ErroringMockPVErr.Error())
+
+			err = tc.mockPV.SignProposal(ctx, tc.chainID, proposal.ToProto())
+			require.Error(t, err)
+
+			err = tc.signerClient.SignProposal(ctx, tc.chainID, proposal.ToProto())
+			require.Error(t, err)
 		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
-			}
-		})
-
-		ts := time.Now()
-		hash := tmrand.Bytes(tmhash.Size)
-		proposal := &types.Proposal{
-			Type:      tmproto.ProposalType,
-			Height:    1,
-			Round:     2,
-			POLRound:  2,
-			BlockID:   types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp: ts,
-			Signature: []byte("signature"),
-		}
-
-		err := tc.signerClient.SignProposal(context.Background(), tc.chainID, proposal.ToProto())
-		require.Equal(t, err.(*RemoteSignerError).Description, types.ErroringMockPVErr.Error())
-
-		err = tc.mockPV.SignProposal(context.Background(), tc.chainID, proposal.ToProto())
-		require.Error(t, err)
-
-		err = tc.signerClient.SignProposal(context.Background(), tc.chainID, proposal.ToProto())
-		require.Error(t, err)
 	}
 }
 
 func TestSignerSignVoteErrors(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		ts := time.Now()
-		hash := tmrand.Bytes(tmhash.Size)
-		valAddr := tmrand.Bytes(crypto.AddressSize)
-		vote := &types.Vote{
-			Type:             tmproto.PrecommitType,
-			Height:           1,
-			Round:            2,
-			BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
-			Timestamp:        ts,
-			ValidatorAddress: valAddr,
-			ValidatorIndex:   1,
-			Signature:        []byte("signature"),
-		}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		// Replace signer service privval with one that always fails
-		tc.signerServer.privVal = types.NewErroringMockPV()
-		tc.mockPV = types.NewErroringMockPV()
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
 
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
+			ts := time.Now()
+			hash := tmrand.Bytes(tmhash.Size)
+			valAddr := tmrand.Bytes(crypto.AddressSize)
+			vote := &types.Vote{
+				Type:             tmproto.PrecommitType,
+				Height:           1,
+				Round:            2,
+				BlockID:          types.BlockID{Hash: hash, PartSetHeader: types.PartSetHeader{Hash: hash, Total: 2}},
+				Timestamp:        ts,
+				ValidatorAddress: valAddr,
+				ValidatorIndex:   1,
+				Signature:        []byte("signature"),
 			}
+
+			// Replace signer service privval with one that always fails
+			tc.signerServer.privVal = types.NewErroringMockPV()
+			tc.mockPV = types.NewErroringMockPV()
+
+			err := tc.signerClient.SignVote(ctx, tc.chainID, vote.ToProto())
+			rserr, ok := err.(*RemoteSignerError)
+			require.True(t, ok, "%T", err)
+			require.Contains(t, rserr.Error(), types.ErroringMockPVErr.Error())
+
+			err = tc.mockPV.SignVote(ctx, tc.chainID, vote.ToProto())
+			require.Error(t, err)
+
+			err = tc.signerClient.SignVote(ctx, tc.chainID, vote.ToProto())
+			require.Error(t, err)
 		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
-			}
-		})
-
-		err := tc.signerClient.SignVote(context.Background(), tc.chainID, vote.ToProto())
-		require.Equal(t, err.(*RemoteSignerError).Description, types.ErroringMockPVErr.Error())
-
-		err = tc.mockPV.SignVote(context.Background(), tc.chainID, vote.ToProto())
-		require.Error(t, err)
-
-		err = tc.signerClient.SignVote(context.Background(), tc.chainID, vote.ToProto())
-		require.Error(t, err)
 	}
 }
 
@@ -413,28 +379,23 @@ func brokenHandler(ctx context.Context, privVal types.PrivValidator, request pri
 }
 
 func TestSignerUnexpectedResponse(t *testing.T) {
-	for _, tc := range getSignerTestCases(t) {
-		tc.signerServer.privVal = types.NewMockPV()
-		tc.mockPV = types.NewMockPV()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-		tc.signerServer.SetRequestHandler(brokenHandler)
+	for _, tc := range getSignerTestCases(ctx, t) {
+		t.Run(tc.name, func(t *testing.T) {
+			defer tc.closer()
 
-		tc := tc
-		t.Cleanup(func() {
-			if err := tc.signerServer.Stop(); err != nil {
-				t.Error(err)
-			}
+			tc.signerServer.privVal = types.NewMockPV()
+			tc.mockPV = types.NewMockPV()
+
+			tc.signerServer.SetRequestHandler(brokenHandler)
+
+			ts := time.Now()
+			want := &types.Vote{Timestamp: ts, Type: tmproto.PrecommitType}
+
+			e := tc.signerClient.SignVote(ctx, tc.chainID, want.ToProto())
+			assert.EqualError(t, e, "empty response")
 		})
-		t.Cleanup(func() {
-			if err := tc.signerClient.Close(); err != nil {
-				t.Error(err)
-			}
-		})
-
-		ts := time.Now()
-		want := &types.Vote{Timestamp: ts, Type: tmproto.PrecommitType}
-
-		e := tc.signerClient.SignVote(context.Background(), tc.chainID, want.ToProto())
-		assert.EqualError(t, e, "empty response")
 	}
 }

privval/signer_listener_endpoint.go

@@ -1,6 +1,7 @@
 package privval
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"time"
@@ -63,7 +64,7 @@ func NewSignerListenerEndpoint(
 }
 
 // OnStart implements service.Service.
-func (sl *SignerListenerEndpoint) OnStart() error {
+func (sl *SignerListenerEndpoint) OnStart(ctx context.Context) error {
 	sl.connectRequestCh = make(chan struct{})
 	sl.connectionAvailableCh = make(chan net.Conn)
 

privval/signer_listener_endpoint_test.go

@@ -1,6 +1,7 @@
 package privval
 
 import (
+	"context"
 	"net"
 	"testing"
 	"time"
@@ -38,6 +39,9 @@ func TestSignerRemoteRetryTCPOnly(t *testing.T) {
 		retries   = 10
 	)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 
@@ -71,7 +75,7 @@ func TestSignerRemoteRetryTCPOnly(t *testing.T) {
 	mockPV := types.NewMockPV()
 	signerServer := NewSignerServer(dialerEndpoint, chainID, mockPV)
 
-	err = signerServer.Start()
+	err = signerServer.Start(ctx)
 	require.NoError(t, err)
 	t.Cleanup(func() {
 		if err := signerServer.Stop(); err != nil {
@@ -88,6 +92,9 @@ func TestSignerRemoteRetryTCPOnly(t *testing.T) {
 }
 
 func TestRetryConnToRemoteSigner(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	for _, tc := range getDialerTestCases(t) {
 		var (
 			logger           = log.TestingLogger()
@@ -107,14 +114,9 @@ func TestRetryConnToRemoteSigner(t *testing.T) {
 
 		signerServer := NewSignerServer(dialerEndpoint, chainID, mockPV)
 
-		startListenerEndpointAsync(t, listenerEndpoint, endpointIsOpenCh)
-		t.Cleanup(func() {
-			if err := listenerEndpoint.Stop(); err != nil {
-				t.Error(err)
-			}
-		})
+		startListenerEndpointAsync(ctx, t, listenerEndpoint, endpointIsOpenCh)
 
-		require.NoError(t, signerServer.Start())
+		require.NoError(t, signerServer.Start(ctx))
 		assert.True(t, signerServer.IsRunning())
 		<-endpointIsOpenCh
 		if err := signerServer.Stop(); err != nil {
@@ -128,13 +130,8 @@ func TestRetryConnToRemoteSigner(t *testing.T) {
 		signerServer2 := NewSignerServer(dialerEndpoint2, chainID, mockPV)
 
 		// let some pings pass
-		require.NoError(t, signerServer2.Start())
+		require.NoError(t, signerServer2.Start(ctx))
 		assert.True(t, signerServer2.IsRunning())
-		t.Cleanup(func() {
-			if err := signerServer2.Stop(); err != nil {
-				t.Error(err)
-			}
-		})
 
 		// give the client some time to re-establish the conn to the remote signer
 		// should see sth like this in the logs:
@@ -175,15 +172,23 @@ func newSignerListenerEndpoint(logger log.Logger, addr string, timeoutReadWrite
 	)
 }
 
-func startListenerEndpointAsync(t *testing.T, sle *SignerListenerEndpoint, endpointIsOpenCh chan struct{}) {
+func startListenerEndpointAsync(
+	ctx context.Context,
+	t *testing.T,
+	sle *SignerListenerEndpoint,
+	endpointIsOpenCh chan struct{},
+) {
+	t.Helper()
+
 	go func(sle *SignerListenerEndpoint) {
-		require.NoError(t, sle.Start())
+		require.NoError(t, sle.Start(ctx))
 		assert.True(t, sle.IsRunning())
 		close(endpointIsOpenCh)
 	}(sle)
 }
 
 func getMockEndpoints(
+	ctx context.Context,
 	t *testing.T,
 	addr string,
 	socketDialer SocketDialer,
@@ -204,9 +209,9 @@ func getMockEndpoints(
 	SignerDialerEndpointTimeoutReadWrite(testTimeoutReadWrite)(dialerEndpoint)
 	SignerDialerEndpointConnRetries(1e6)(dialerEndpoint)
 
-	startListenerEndpointAsync(t, listenerEndpoint, endpointIsOpenCh)
+	startListenerEndpointAsync(ctx, t, listenerEndpoint, endpointIsOpenCh)
 
-	require.NoError(t, dialerEndpoint.Start())
+	require.NoError(t, dialerEndpoint.Start(ctx))
 	assert.True(t, dialerEndpoint.IsRunning())
 
 	<-endpointIsOpenCh

privval/signer_server.go

@@ -42,8 +42,8 @@ func NewSignerServer(endpoint *SignerDialerEndpoint, chainID string, privVal typ
 }
 
 // OnStart implements service.Service.
-func (ss *SignerServer) OnStart() error {
-	go ss.serviceLoop()
+func (ss *SignerServer) OnStart(ctx context.Context) error {
+	go ss.serviceLoop(ctx)
 	return nil
 }
 
@@ -91,18 +91,18 @@ func (ss *SignerServer) servicePendingRequest() {
 	}
 }
 
-func (ss *SignerServer) serviceLoop() {
+func (ss *SignerServer) serviceLoop(ctx context.Context) {
 	for {
 		select {
+		case <-ss.Quit():
+			return
+		case <-ctx.Done():
+			return
 		default:
-			err := ss.endpoint.ensureConnection()
-			if err != nil {
+			if err := ss.endpoint.ensureConnection(); err != nil {
 				return
 			}
 			ss.servicePendingRequest()
-
-		case <-ss.Quit():
-			return
 		}
 	}
 }