test: Improve TestSSH/ForwardGPG stability on macOS via pty.ReadRune (#5739)

mafredri · web-flow · commit 77e71f3ca48c · 2023-01-17T15:30:47.000+02:00
Writing to stdin for `coder ssh` too early could result in the input
being discarded. To work around this we add a new `ptytest` method
called `ReadRune` that lets us read one character of output. This will
indicate the command is ready to accept input.

It could be one character of the prompt, or of the loading message
waiting for connection to be established.
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -475,6 +475,10 @@ Expire-Date: 0
 		// real error from being printed.
 		t.Cleanup(cancel)
 
+		// Wait for the prompt or any output really to indicate the command has
+		// started and accepting input on stdin.
+		_ = pty.ReadRune(ctx)
+
 		pty.WriteLine("echo hello 'world'")
 		pty.ExpectMatch("hello world")
 
diff --git a/pty/ptytest/ptytest.go b/pty/ptytest/ptytest.go
@@ -182,6 +182,47 @@ func (p *PTY) ExpectMatch(str string) string {
 	}
 }
 
+func (p *PTY) ReadRune(ctx context.Context) rune {
+	p.t.Helper()
+
+	// A timeout is mandatory, caller can decide by passing a context
+	// that times out.
+	if _, ok := ctx.Deadline(); !ok {
+		timeout := testutil.WaitMedium
+		p.logf("ReadRune ctx has no deadline, using %s", timeout)
+		var cancel context.CancelFunc
+		//nolint:gocritic // Rule guard doesn't detect that we're using testutil.Wait*.
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
+
+	var r rune
+	match := make(chan error, 1)
+	go func() {
+		defer close(match)
+		var err error
+		r, _, err = p.runeReader.ReadRune()
+		match <- err
+	}()
+
+	select {
+	case err := <-match:
+		if err != nil {
+			p.fatalf("read error", "%v (wanted newline; got %q)", err, r)
+			return 0
+		}
+		p.logf("matched rune = %q", r)
+		return r
+	case <-ctx.Done():
+		// Ensure goroutine is cleaned up before test exit.
+		_ = p.close("read rune context done: " + ctx.Err().Error())
+		<-match
+
+		p.fatalf("read rune context done", "wanted rune; got nothing")
+		return 0
+	}
+}
+
 func (p *PTY) ReadLine() string {
 	p.t.Helper()
 
