diff --git a/internal/cli/ws/dial_test.go b/internal/cli/ws/dial_test.go index 364533c..5b40568 100644 --- a/internal/cli/ws/dial_test.go +++ b/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) - hub.Push("event", map[string]any{"n": 1}) - - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - hub.Shutdown(ctx) + done <- dial(context.Background(), nil, out, dialSettings{ + url: url, + messages: []string{ping}, + listen: true, + timeout: time.Second, + }) }() - 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)", - ) + // 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) + + 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