internal/master/hub/jobs_test.go

@@ -2,6 +2,7 @@ package hub
 
 import (
 	"context"
+	"errors"
 	"strconv"
 	"strings"
 	"sync"
@@ -106,9 +107,12 @@ func TestHub_StartJobOnReplacedConnectionFails(t *testing.T) {
 	require.NoError(t, err)
 	defer second.CloseNow()
 
-	// Until the new connection says hello, nothing is routable and the old
-	// route is refused.
-	assert.ErrorIs(t, h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)
+	// dial returns before the hub has registered the connection, so wait for
+	// the replacement to take effect: the old route is then refused. Until the
+	// new connection says hello, nothing is routable.
+	require.Eventually(t, func() bool {
+		return errors.Is(h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)
+	}, 2*time.Second, 5*time.Millisecond)
 	send(t, second, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
 	fresh := waitRoute(t, h)
 	require.NotEqual(t, old.ConnID, fresh.ConnID)