Fixed a race condition in dial_test related to hub closing in-test

internal/cli/ws/dial_test.go

@@ -161,32 +161,38 @@ func TestDialOrigin(t *testing.T) {
 // With --listen, pushed messages are printed until the server closes the
 // connection.
 func TestDialListen(t *testing.T) {
-	cmdtest.Setup(t, "")
 	hub, url := wsServer(t)
+	out := &syncBuffer{}
 
+	done := make(chan error, 1)
 	go func() {
-		// Wait for the client to connect and get its pong.
-		deadline := time.Now().Add(2 * time.Second)
-		for hub.Count() == 0 && time.Now().Before(deadline) {
-			time.Sleep(5 * time.Millisecond)
-		}
-		time.Sleep(50 * time.Millisecond)
+		done <- dial(context.Background(), nil, out, dialSettings{
+			url:      url,
+			messages: []string{ping},
+			listen:   true,
+			timeout:  time.Second,
+		})
+	}()
+
+	// Push only once the client has its pong, and shut down only once it has
+	// the push: Shutdown closes clients without waiting for their queues.
+	waitFor(t, out, `< {"type":"pong","id":"1"`)
 	hub.Push("event", map[string]any{"n": 1})
+	waitFor(t, out, `< {"type":"event","payload":{"n":1}}`)
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 	hub.Shutdown(ctx)
-	}()
 
-	out := cmdtest.MustRun(
-		t, newRoot, "ws", "dial", "--url", url, "--message", ping, "--listen",
-	)
-	cmdtest.Contains(
-		t,
-		out,
-		`< {"type":"event","payload":{"n":1}}`,
-		"closed by the server (1001)",
-	)
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("dial = %v; want nil when the server closes", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("dial did not return when the server closed")
+	}
+	waitFor(t, out, "closed by the server (1001)")
 }
 
 // syncBuffer is a bytes.Buffer safe to write from one goroutine while