[deliver] Fix flaky test

Signed-off-by: Liran Funaro <liran.funaro@gmail.com>

Author: liran-funaro

utils/deliver/client_test.go

@@ -51,9 +51,7 @@ func TestBroadcastDeliver(t *testing.T) {
 
 	ctx, cancel := context.WithTimeout(t.Context(), time.Minute)
 	t.Cleanup(cancel)
-	stream, err := test.NewBroadcastStream(ctx, &conf)
-	require.NoError(t, err)
-	t.Cleanup(stream.Close)
+
 	outputBlocksChan := make(chan *common.Block, 100)
 	go func() {
 		p := &Parameters{
@@ -78,59 +76,62 @@ func TestBroadcastDeliver(t *testing.T) {
 	require.Len(t, b.Data.Data, 1)
 
 	t.Log("All good")
-	submit(t, stream, outputBlocks, expectedSubmit{
+	submit(t, &conf, outputBlocks, expectedSubmit{
 		success: 3,
 	})
 
 	t.Log("One server down")
 	servers[2].Servers[0].Stop()
-	time.Sleep(3 * time.Second)
-	submit(t, stream, outputBlocks, expectedSubmit{
+	waitUntilGrpcServerIsDown(ctx, t, &servers[2].Configs[0].Endpoint)
+	submit(t, &conf, outputBlocks, expectedSubmit{
 		success: 3,
 	})
 
 	t.Log("Two servers down")
 	servers[2].Servers[1].Stop()
-	time.Sleep(3 * time.Second)
-	submit(t, stream, outputBlocks, expectedSubmit{
+	waitUntilGrpcServerIsDown(ctx, t, &servers[2].Configs[1].Endpoint)
+	submit(t, &conf, outputBlocks, expectedSubmit{
 		success:     2,
 		unavailable: 1,
 	})
 
 	t.Log("One incorrect server")
 	fakeServer := test.RunGrpcServerForTest(ctx, t, servers[2].Configs[0], nil)
 	waitUntilGrpcServerIsReady(ctx, t, &servers[2].Configs[0].Endpoint)
-
-	submit(t, stream, outputBlocks, expectedSubmit{
+	submit(t, &conf, outputBlocks, expectedSubmit{
 		success:       2,
 		unimplemented: 1,
 	})
 
 	t.Log("All good again")
 	fakeServer.Stop()
+	waitUntilGrpcServerIsDown(ctx, t, &servers[2].Configs[0].Endpoint)
 	servers[2].Servers[0] = test.RunGrpcServerForTest(ctx, t, servers[2].Configs[0], ordererService.RegisterService)
 	waitUntilGrpcServerIsReady(ctx, t, &servers[2].Configs[0].Endpoint)
-	submit(t, stream, outputBlocks, expectedSubmit{
+	submit(t, &conf, outputBlocks, expectedSubmit{
 		success: 3,
 	})
 
 	t.Log("Insufficient quorum")
-	servers[0].Servers[0].Stop()
-	servers[0].Servers[1].Stop()
-	servers[1].Servers[0].Stop()
-	servers[1].Servers[1].Stop()
-	time.Sleep(3 * time.Second)
-	submit(t, stream, outputBlocks, expectedSubmit{
+	for i := range 2 {
+		for j := range 2 {
+			servers[i].Servers[j].Stop()
+		}
+	}
+	for i := range 2 {
+		for j := range 2 {
+			waitUntilGrpcServerIsDown(ctx, t, &servers[i].Configs[j].Endpoint)
+		}
+	}
+	submit(t, &conf, outputBlocks, expectedSubmit{
 		success:     1,
 		unavailable: 2,
 	})
 
 	t.Log("Update endpoints")
 	conf.Connection.Endpoints = allEndpoints[6:]
-	err = client.UpdateConnections(&conf.Connection)
-	require.NoError(t, err)
-	require.NoError(t, stream.ConnectionManager.Update(&conf.Connection))
-	submit(t, stream, outputBlocks, expectedSubmit{
+	require.NoError(t, client.UpdateConnections(&conf.Connection))
+	submit(t, &conf, outputBlocks, expectedSubmit{
 		success: 3,
 	})
 }
@@ -143,22 +144,29 @@ type expectedSubmit struct {
 
 func submit(
 	t *testing.T,
-	stream *test.BroadcastStream,
+	conf *ordererconn.Config,
 	outputBlocks channel.Reader[*common.Block],
 	expected expectedSubmit,
 ) {
 	t.Helper()
 	txb := workload.TxBuilder{ChannelID: channelForTest}
 	tx := txb.MakeTx(&protoblocktx.Tx{})
 
-	err := stream.SendBatch(workload.MapToEnvelopeBatch(0, []*protoloadgen.TX{tx}))
+	// We create a new stream for each request to ensure GRPC does not cache the latest state.
+	ctx, cancel := context.WithTimeout(t.Context(), time.Minute)
+	defer cancel()
+	stream, err := test.NewBroadcastStream(ctx, conf)
+	require.NoError(t, err)
+	defer stream.Close()
+
+	err = stream.SendBatch(workload.MapToEnvelopeBatch(0, []*protoloadgen.TX{tx}))
 	if err != nil {
 		t.Logf("Response error:\n%s", err)
 	}
 	if expected.unavailable > 0 {
-		// Unimplemented sometimes not responding with error, so we cannot enforce it.
 		require.Error(t, err)
 	} else if expected.unimplemented == 0 {
+		// Unimplemented sometimes not responding with error, so we cannot enforce it.
 		require.NoError(t, err)
 	}
 
@@ -226,6 +234,13 @@ func waitUntilGrpcServerIsReady(ctx context.Context, t *testing.T, endpoint *con
 	defer connection.CloseConnectionsLog(newConn)
 	test.WaitUntilGrpcServerIsReady(ctx, t, newConn)
 	t.Logf("%v is ready", endpoint)
-	// Wait a while to allow the GRPC connection enough time to make a new attempt.
-	time.Sleep(time.Second)
+}
+
+func waitUntilGrpcServerIsDown(ctx context.Context, t *testing.T, endpoint *connection.Endpoint) {
+	t.Helper()
+	newConn, err := connection.Connect(connection.NewInsecureDialConfig(endpoint))
+	require.NoError(t, err)
+	defer connection.CloseConnectionsLog(newConn)
+	test.WaitUntilGrpcServerIsDown(ctx, t, newConn)
+	t.Logf("%v is down", endpoint)
 }

utils/test/utils.go

@@ -206,12 +206,31 @@ func WaitUntilGrpcServerIsReady(
 		return
 	}
 	healthClient := healthgrpc.NewHealthClient(conn)
-	res, err := healthClient.Check(ctx, &healthgrpc.HealthCheckRequest{}, grpc.WaitForReady(true))
+	res, err := healthClient.Check(ctx, nil, grpc.WaitForReady(true))
 	assert.NotEqual(t, codes.Canceled, status.Code(err))
 	require.NoError(t, err)
 	require.Equal(t, healthgrpc.HealthCheckResponse_SERVING, res.Status)
 }
 
+// WaitUntilGrpcServerIsDown uses the health check API to check a service is down.
+func WaitUntilGrpcServerIsDown(
+	ctx context.Context,
+	t *testing.T,
+	conn grpc.ClientConnInterface,
+) {
+	t.Helper()
+	if conn == nil {
+		return
+	}
+	healthClient := healthgrpc.NewHealthClient(conn)
+	require.EventuallyWithT(t, func(ct *assert.CollectT) {
+		checkCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+		defer cancel()
+		_, err := healthClient.Check(checkCtx, nil)
+		require.Error(ct, err)
+	}, time.Minute, 50*time.Millisecond)
+}
+
 // StatusRetriever provides implementation retrieve status of given transaction identifiers.
 type StatusRetriever interface {
 	GetTransactionsStatus(context.Context, *protoblocktx.QueryStatus, ...grpc.CallOption) (