[Testing] Restore behavioral bounds weakened by flakiness fixes

engine/access/rest/websockets/controller_test.go

@@ -888,6 +888,11 @@ func (s *WsControllerSuite) TestRateLimiter() {
 	elapsed := timestamps[len(timestamps)-1].Sub(timestamps[0])
 	assert.GreaterOrEqual(t, elapsed, minElapsed-tolerance,
 		"Messages should be paced at no more than the configured rate")
+	// Generous upper bound: catches a limiter that stalls the stream. Emitting the messages
+	// must not take dramatically longer than the rate limit dictates; the slack absorbs
+	// scheduling delays on a loaded machine without reintroducing flakiness.
+	assert.LessOrEqual(t, elapsed, minElapsed+2*time.Second,
+		"Messages should not be excessively delayed by the rate limiter")
 }
 
 // TestConfigureKeepaliveConnection ensures that the WebSocket connection is configured correctly.

engine/access/rpc/connection/connection_test.go

@@ -270,10 +270,14 @@ func TestExecutionNodeClientTimeout(t *testing.T) {
 	// setup the handler mock to not respond within the timeout
 	req := &execution.PingRequest{}
 	resp := &execution.PingResponse{}
+	handlerReached := atomic.NewInt64(0)
 	en.handler.
 		On("Ping",
 			testifymock.Anything,
 			testifymock.AnythingOfType("*execution.PingRequest")).
+		Run(func(_ testifymock.Arguments) {
+			handlerReached.Inc()
+		}).
 		After(timeout+time.Second).
 		Return(resp, nil).
 		Maybe() // under load, the client's deadline can fire before the request even reaches the server
@@ -305,11 +309,15 @@ func TestExecutionNodeClientTimeout(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx := context.Background()
-	// make the call to the execution node
-	_, err = client.Ping(ctx, req)
-
-	// assert that the client timed out
-	assert.Equal(t, codes.DeadlineExceeded, status.Code(err))
+	// Make the call to the execution node. Every attempt must time out. Under load, an attempt
+	// can time out before the request even reaches the server (e.g. during connection setup);
+	// retry a few times so that at least one timeout is deterministically caused by the slow
+	// handler, which is the behavior under test.
+	require.Eventually(t, func() bool {
+		_, err = client.Ping(ctx, req)
+		assert.Equal(t, codes.DeadlineExceeded, status.Code(err))
+		return handlerReached.Load() > 0
+	}, 5*time.Second, 10*time.Millisecond, "server handler was never reached; timeouts were not caused by the slow handler")
 }
 
 // TestCollectionNodeClientTimeout tests that the collection API client times out after the timeout duration
@@ -327,10 +335,14 @@ func TestCollectionNodeClientTimeout(t *testing.T) {
 	// setup the handler mock to not respond within the timeout
 	req := &access.PingRequest{}
 	resp := &access.PingResponse{}
+	handlerReached := atomic.NewInt64(0)
 	cn.handler.
 		On("Ping",
 			testifymock.Anything,
 			testifymock.AnythingOfType("*access.PingRequest")).
+		Run(func(_ testifymock.Arguments) {
+			handlerReached.Inc()
+		}).
 		After(timeout+time.Second).
 		Return(resp, nil).
 		Maybe() // under load, the client's deadline can fire before the request even reaches the server
@@ -362,11 +374,15 @@ func TestCollectionNodeClientTimeout(t *testing.T) {
 	assert.NoError(t, err)
 
 	ctx := context.Background()
-	// make the call to the execution node
-	_, err = client.Ping(ctx, req)
-
-	// assert that the client timed out
-	assert.Equal(t, codes.DeadlineExceeded, status.Code(err))
+	// Make the call to the collection node. Every attempt must time out. Under load, an attempt
+	// can time out before the request even reaches the server (e.g. during connection setup);
+	// retry a few times so that at least one timeout is deterministically caused by the slow
+	// handler, which is the behavior under test.
+	require.Eventually(t, func() bool {
+		_, err = client.Ping(ctx, req)
+		assert.Equal(t, codes.DeadlineExceeded, status.Code(err))
+		return handlerReached.Load() > 0
+	}, 5*time.Second, 10*time.Millisecond, "server handler was never reached; timeouts were not caused by the slow handler")
 }
 
 // TestConnectionPoolFull tests that the LRU cache replaces connections when full

network/alsp/manager/manager_test.go

@@ -1595,6 +1595,7 @@ func TestDecayMisbehaviorPenalty_MultipleHeartbeats(t *testing.T) {
 	// fewer heartbeats (flaky). Instead, we poll until 3 heartbeats worth of decay is visible. The
 	// poll interval (10ms) is much smaller than the heartbeat interval (1s), so we observe the
 	// record right after the third decaying heartbeat, before a fourth one can fire.
+	decayStart := time.Now()
 	var record *model.ProtocolSpamRecord
 	require.Eventually(t, func() bool {
 		r, ok := cache.Get(originId)
@@ -1609,6 +1610,13 @@ func TestDecayMisbehaviorPenalty_MultipleHeartbeats(t *testing.T) {
 		return true
 	}, 10*time.Second, 10*time.Millisecond, "penalty was not decayed by 3 heartbeats on time")
 
+	// Pacing: decays happen at most once per heartbeat interval. Reaching 3 heartbeats worth of
+	// decay requires at least 2 full intervals after the first decaying heartbeat (which may
+	// fire immediately after the record was created). This is a lower bound only, so it cannot
+	// flake on a slow machine.
+	require.GreaterOrEqual(t, time.Since(decayStart), 2*cfg.HeartBeatInterval,
+		"penalty decayed faster than the heartbeat interval allows")
+
 	// observed right after the third decaying heartbeat, the penalty must be less than the value after 4 heartbeats.
 	require.Less(t, record.Penalty, penaltyBeforeDecay+4*record.Decay)
 	require.False(t, record.DisallowListed) // the peer should not be disallow listed yet.

utils/unittest/unittest.go

@@ -317,9 +317,14 @@ func TempDir(t testing.TB) string {
 // appear in the directory while os.RemoveAll is deleting it, failing it with ENOTEMPTY.
 func RemoveTempDir(t testing.TB, dir string) {
 	var err error
-	for deadline := time.Now().Add(10 * time.Second); time.Now().Before(deadline); {
+	for attempt, deadline := 0, time.Now().Add(10*time.Second); time.Now().Before(deadline); attempt++ {
 		err = os.RemoveAll(dir)
 		if err == nil {
+			if attempt > 0 {
+				// preserve the signal that something was still writing to the directory: this
+				// usually means a component leaked a background writer past its shutdown
+				t.Logf("removing temp dir %s required %d retries, something was still writing to it", dir, attempt)
+			}
 			return
 		}
 		time.Sleep(100 * time.Millisecond)