test(lsp): wait on observable conditions instead of fixed sleeps

cmd/lsp_stdio_test.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 )
@@ -49,6 +50,84 @@ func buildBinary(t *testing.T) string {
 	return bin
 }
 
+// lockedBuffer collects a child's output while the test reads it.
+type lockedBuffer struct {
+	mu sync.Mutex
+	b  bytes.Buffer
+}
+
+func (l *lockedBuffer) Write(p []byte) (int, error) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.b.Write(p)
+}
+
+func (l *lockedBuffer) Bytes() []byte {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return bytes.Clone(l.b.Bytes())
+}
+
+func (l *lockedBuffer) String() string { return string(l.Bytes()) }
+
+func writeAll(t *testing.T, w io.Writer, msgs [][]byte) {
+	t.Helper()
+	for _, m := range msgs {
+		if _, err := w.Write(m); err != nil {
+			t.Fatalf("writing to the server: %v", err)
+		}
+	}
+}
+
+// waitForMessage polls the server's stdout until a complete message matches.
+// A frame still being written is read on a later poll; whether the stream as a
+// whole is clean framing is the strict parse at the end of the test.
+func waitForMessage(t *testing.T, out *lockedBuffer, timeout time.Duration, what string, match func(map[string]any) bool) map[string]any {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for {
+		frames, _ := parseFrames(bytes.NewReader(out.Bytes()))
+		for _, f := range frames {
+			var m map[string]any
+			if json.Unmarshal(f, &m) == nil && match(m) {
+				return m
+			}
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("no %s on stdout within %s", what, timeout)
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+}
+
+// isResponse matches the response to the request with the given id.
+func isResponse(id int) func(map[string]any) bool {
+	return func(m map[string]any) bool {
+		_, hasMethod := m["method"]
+		return !hasMethod && fmt.Sprint(m["id"]) == strconv.Itoa(id)
+	}
+}
+
+// waitForRules asks vulnetix/serverInfo until it reports the loaded rules,
+// which the server stores once the warm-up has indexed the workspace.
+func waitForRules(t *testing.T, stdin io.Writer, out *lockedBuffer, timeout time.Duration) {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for id := 100; ; id++ {
+		writeAll(t, stdin, [][]byte{frame(t, map[string]any{"jsonrpc": "2.0", "id": id, "method": "vulnetix/serverInfo"})})
+		resp := waitForMessage(t, out, timeout, "a serverInfo response", isResponse(id))
+		if result, ok := resp["result"].(map[string]any); ok {
+			if n, _ := result["rulesEmbedded"].(float64); n > 0 {
+				return
+			}
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("the server did not load its rules within %s", timeout)
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+}
+
 // frame renders a JSON-RPC message with LSP framing.
 func frame(t *testing.T, v any) []byte {
 	t.Helper()
@@ -136,7 +215,7 @@ func TestNoStdoutPollution(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	var stdout, stderr bytes.Buffer
+	var stdout, stderr lockedBuffer
 	cmd.Stdout = &stdout
 	cmd.Stderr = &stderr
 
@@ -164,6 +243,14 @@ func TestNoStdoutPollution(t *testing.T) {
 		"jsonrpc": "2.0", "method": "initialized", "params": map[string]any{},
 	}))
 
+	writeAll(t, stdin, msgs)
+	msgs = nil
+
+	// A document opened before the warm-up has loaded the rules is not analysed
+	// until its next change, so wait for serverInfo to report them before
+	// opening one. These requests are protocol on stdout like any other.
+	waitForRules(t, stdin, &stdout, 2*time.Minute)
+
 	docURI := "file://" + filepath.ToSlash(filepath.Join(root, "main.go"))
 	msgs = append(msgs, frame(t, map[string]any{
 		"jsonrpc": "2.0", "method": "textDocument/didOpen",
@@ -184,18 +271,17 @@ func TestNoStdoutPollution(t *testing.T) {
 		"params": map[string]any{"textDocument": map[string]any{"uri": docURI}},
 	}))
 
-	for _, m := range msgs {
-		if _, err := stdin.Write(m); err != nil {
-			t.Fatalf("writing to the server: %v", err)
-		}
-	}
+	writeAll(t, stdin, msgs)
 
-	// Give analysis a chance to run and publish, which is when a stray write is
-	// most likely.
-	time.Sleep(6 * time.Second)
+	// Wait for analysis to run and publish, which is when a stray write is most
+	// likely.
+	waitForMessage(t, &stdout, 2*time.Minute, "diagnostics for "+docURI, func(m map[string]any) bool {
+		params, _ := m["params"].(map[string]any)
+		return m["method"] == "textDocument/publishDiagnostics" && params["uri"] == docURI
+	})
 
 	_, _ = stdin.Write(frame(t, map[string]any{"jsonrpc": "2.0", "id": 2, "method": "shutdown"}))
-	time.Sleep(300 * time.Millisecond)
+	waitForMessage(t, &stdout, 15*time.Second, "the shutdown response", isResponse(2))
 	_, _ = stdin.Write(frame(t, map[string]any{"jsonrpc": "2.0", "method": "exit"}))
 	_ = stdin.Close()
 

internal/lsp/conformance_test.go

@@ -248,7 +248,7 @@ func TestDidOpenPublishesDiagnostics(t *testing.T) {
 	c, stop := startServer(t, root)
 	defer stop()
 	initialize(t, c, root)
-	waitForWarm(t)
+	waitForWarm(t, c)
 
 	uri := PathToURI(vulnPath)
 	require.NoError(t, c.conn.Notify(context.Background(), protocol.MethodDidOpen, protocol.DidOpenParams{
@@ -279,7 +279,7 @@ func TestClosingADocumentClearsItsDiagnostics(t *testing.T) {
 	c, stop := startServer(t, root)
 	defer stop()
 	initialize(t, c, root)
-	waitForWarm(t)
+	waitForWarm(t, c)
 
 	uri := PathToURI(vulnPath)
 	require.NoError(t, c.conn.Notify(context.Background(), protocol.MethodDidOpen, protocol.DidOpenParams{
@@ -325,14 +325,25 @@ func TestDocumentsOutsideTheWorkspaceAreNotScanned(t *testing.T) {
 	c, stop := startServer(t, root)
 	defer stop()
 	initialize(t, c, root)
-	waitForWarm(t)
+	waitForWarm(t, c)
 
 	uri := PathToURI(outside)
 	require.NoError(t, c.conn.Notify(context.Background(), protocol.MethodDidOpen, protocol.DidOpenParams{
 		TextDocument: protocol.TextDocumentItem{URI: uri, LanguageID: "go", Version: 1, Text: vulnerableGo},
 	}))
 
-	time.Sleep(500 * time.Millisecond)
+	// Barrier: open a file inside the workspace after the outside one and wait
+	// for its diagnostics. The outside file is turned away when it is opened,
+	// while the inside one has to be analysed, so by the time the inside one is
+	// published a publish for the outside one would have arrived too.
+	inside := filepath.Join(root, "vuln.go")
+	require.NoError(t, os.WriteFile(inside, []byte(vulnerableGo), 0o644))
+	insideURI := PathToURI(inside)
+	require.NoError(t, c.conn.Notify(context.Background(), protocol.MethodDidOpen, protocol.DidOpenParams{
+		TextDocument: protocol.TextDocumentItem{URI: insideURI, LanguageID: "go", Version: 1, Text: vulnerableGo},
+	}))
+	require.NotEmpty(t, c.waitForPublish(t, insideURI, 2*time.Minute), "the in-workspace file should be analysed")
+
 	require.Empty(t, c.diagsFor(uri), "a file outside the workspace should not be analysed")
 }
 
@@ -364,15 +375,25 @@ func TestWorkspaceScanCanBeCancelled(t *testing.T) {
 
 // ── Helpers ──────────────────────────────────────────────────────────────────
 
-// waitForWarm gives the background warm-up time to index and compile.
+// waitForWarm waits until the background warm-up has loaded the rules and
+// indexed the workspace.
 //
 // The server deliberately returns from initialize before compiling, so the
-// editor window is not blocked. Tests that assert on findings have to wait for
-// it, and a generous ceiling costs nothing when the condition is usually met
-// in well under a second.
-func waitForWarm(t *testing.T) {
+// editor window is not blocked, and a document opened before the warm-up ends
+// is not analysed until its next change. serverInfo reports the loaded rules
+// once the warm-up has stored them together with the workspace index, which is
+// the point at which a didOpen is analysed. Polling it returns as soon as that
+// happens instead of sleeping a fixed time that is too long on a clean run and
+// too short on a loaded one.
+func waitForWarm(t *testing.T, c *testClient) {
 	t.Helper()
-	time.Sleep(3 * time.Second)
+	require.Eventually(t, func() bool {
+		var info protocol.ServerInfoResult
+		if err := c.conn.Call(context.Background(), protocol.MethodServerInfo, nil, &info); err != nil {
+			return false
+		}
+		return info.RulesEmbedded > 0
+	}, 2*time.Minute, 20*time.Millisecond, "the server never finished warming up")
 }
 
 const vulnerableGo = `package main