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.
+		done <- dial(context.Background(), nil, out, dialSettings{
-		deadline := time.Now().Add(2 * time.Second)
+			url:      url,
-		for hub.Count() == 0 && time.Now().Before(deadline) {
+			messages: []string{ping},
-			time.Sleep(5 * time.Millisecond)
+			listen:   true,
-		}
+			timeout:  time.Second,
-		time.Sleep(50 * time.Millisecond)
+		})
-		hub.Push("event", map[string]any{"n": 1})
-
-		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-		defer cancel()
-		hub.Shutdown(ctx)
 	}()
 
-	out := cmdtest.MustRun(
+	// Push only once the client has its pong, and shut down only once it has
-		t, newRoot, "ws", "dial", "--url", url, "--message", ping, "--listen",
+	// the push: Shutdown closes clients without waiting for their queues.
-	)
+	waitFor(t, out, `< {"type":"pong","id":"1"`)
-	cmdtest.Contains(
+	hub.Push("event", map[string]any{"n": 1})
-		t,
+	waitFor(t, out, `< {"type":"event","payload":{"n":1}}`)
-		out,
+
-		`< {"type":"event","payload":{"n":1}}`,
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-		"closed by the server (1001)",
+	defer cancel()
-	)
+	hub.Shutdown(ctx)
+
+	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