tweak test

internal/consensus/invalid_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -20,7 +21,7 @@ import (
 )
 
 func TestReactorInvalidPrecommit(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
 	config := configSetup(t)
@@ -49,14 +50,14 @@ func TestReactorInvalidPrecommit(t *testing.T) {
 	byzState := rts.states[node.NodeID]
 	byzReactor := rts.reactors[node.NodeID]
 
-	calledDoPrevote := false
+	signal := make(chan struct{})
 	// Update the doPrevote function to just send a valid precommit for a random
 	// block and otherwise disable the priv validator.
 	byzState.mtx.Lock()
 	privVal := byzState.privValidator
 	byzState.doPrevote = func(ctx context.Context, height int64, round int32) {
+		defer close(signal)
 		invalidDoPrevoteFunc(ctx, t, height, round, byzState, byzReactor, privVal)
-		calledDoPrevote = true
 	}
 	byzState.mtx.Unlock()
 
@@ -72,16 +73,30 @@ func TestReactorInvalidPrecommit(t *testing.T) {
 			go func(s eventbus.Subscription) {
 				defer wg.Done()
 				_, err := s.Next(ctx)
+				if ctx.Err() != nil {
+					return
+				}
 				if !assert.NoError(t, err) {
 					cancel() // cancel other subscribers on failure
 				}
 			}(sub)
 		}
 	}
+	wait := make(chan struct{})
+	go func() { defer close(wait); wg.Wait() }()
 
-	wg.Wait()
-	if !calledDoPrevote {
-		t.Fatal("test failed to run core logic")
+	select {
+	case <-wait:
+		if _, ok := <-signal; !ok {
+			t.Fatal("test condition did not fire")
+		}
+	case <-ctx.Done():
+		if _, ok := <-signal; !ok {
+			t.Fatal("test condition did not fire after timeout")
+			return
+		}
+	case <-signal:
+		// test passed
 	}
 }
 
@@ -130,19 +145,27 @@ func invalidDoPrevoteFunc(
 		cs.privValidator = nil // disable priv val so we don't do normal votes
 		cs.mtx.Unlock()
 
-		count := 0
+		r.mtx.Lock()
+		ids := make([]types.NodeID, 0, len(r.peers))
 		for _, ps := range r.peers {
+			ids = append(ids, ps.peerID)
+		}
+		r.mtx.Unlock()
+
+		count := 0
+		for _, peerID := range ids {
 			count++
 			err := r.voteCh.Send(ctx, p2p.Envelope{
-				To: ps.peerID,
+				To: peerID,
 				Message: &tmcons.Vote{
 					Vote: precommit.ToProto(),
 				},
 			})
 			// we want to have sent some of these votes,
 			// but if the test completes without erroring
-			// and we get here, we shouldn't error
-			if errors.Is(err, context.Canceled) && count > 1 {
+			// or not sending any messages, then we should
+			// error.
+			if errors.Is(err, context.Canceled) && count > 0 {
 				break
 			}
 			require.NoError(t, err)