internal/lsp/regtest: track outstanding work using the progress API

In preparation for later changes, add support for tracking outstanding
work in the lsp regtests. This simply threads through progress
notifications and tracks their state in regtest.Env.state. A new
Expectation is added to assert that there is no outstanding work, but
this is as-yet unused.

A unit test is added for Env to check that we're handling work progress
reports correctly after Marshaling/Unmarshaling, since we're not yet
exercising this code path in actual regtests.

Change-Id: I104caf25cfd49340f13d086314f5aef2b8f3bd3b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/229320
Run-TryBot: Robert Findley <rfindley@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/regtest/env.go b/internal/lsp/regtest/env.go
index 486beb0..2bec2c5 100644
--- a/internal/lsp/regtest/env.go
+++ b/internal/lsp/regtest/env.go
@@ -346,6 +346,15 @@
 	// diagnostics are a map of relative path->diagnostics params
 	diagnostics map[string]*protocol.PublishDiagnosticsParams
 	logs        []*protocol.LogMessageParams
+	// outstandingWork is a map of token->work summary. All tokens are assumed to
+	// be string, though the spec allows for numeric tokens as well.  When work
+	// completes, it is deleted from this map.
+	outstandingWork map[string]*workProgress
+}
+
+type workProgress struct {
+	title   string
+	percent float64
 }
 
 func (s State) String() string {
@@ -368,6 +377,15 @@
 			fmt.Fprintf(&b, "\t\t(%d, %d): %s\n", int(d.Range.Start.Line), int(d.Range.Start.Character), d.Message)
 		}
 	}
+	b.WriteString("\n")
+	b.WriteString("#### outstanding work:\n")
+	for token, state := range s.outstandingWork {
+		name := state.title
+		if name == "" {
+			name = fmt.Sprintf("!NO NAME(token: %s)", token)
+		}
+		fmt.Fprintf(&b, "\t%s: %.2f", name, state.percent)
+	}
 	return b.String()
 }
 
@@ -396,12 +414,15 @@
 		Server: ts,
 		Conn:   conn,
 		state: State{
-			diagnostics: make(map[string]*protocol.PublishDiagnosticsParams),
+			diagnostics:     make(map[string]*protocol.PublishDiagnosticsParams),
+			outstandingWork: make(map[string]*workProgress),
 		},
 		waiters: make(map[int]*condition),
 	}
 	env.E.Client().OnDiagnostics(env.onDiagnostics)
 	env.E.Client().OnLogMessage(env.onLogMessage)
+	env.E.Client().OnWorkDoneProgressCreate(env.onWorkDoneProgressCreate)
+	env.E.Client().OnProgress(env.onProgress)
 	return env
 }
 
@@ -423,6 +444,38 @@
 	return nil
 }
 
+func (e *Env) onWorkDoneProgressCreate(_ context.Context, m *protocol.WorkDoneProgressCreateParams) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	// panic if we don't have a string token.
+	token := m.Token.(string)
+	e.state.outstandingWork[token] = &workProgress{}
+	return nil
+}
+
+func (e *Env) onProgress(_ context.Context, m *protocol.ProgressParams) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	token := m.Token.(string)
+	work, ok := e.state.outstandingWork[token]
+	if !ok {
+		panic(fmt.Sprintf("got progress report for unknown report %s", token))
+	}
+	v := m.Value.(map[string]interface{})
+	switch kind := v["kind"]; kind {
+	case "begin":
+		work.title = v["title"].(string)
+	case "report":
+		if pct, ok := v["percentage"]; ok {
+			work.percent = pct.(float64)
+		}
+	case "end":
+		delete(e.state.outstandingWork, token)
+	}
+	e.checkConditionsLocked()
+	return nil
+}
+
 func (e *Env) checkConditionsLocked() {
 	for id, condition := range e.waiters {
 		if v, _, _ := checkExpectations(e.state, condition.expectations); v != Unmet {
@@ -512,6 +565,37 @@
 	return fmt.Sprintf("unrecognized verdict %d", v)
 }
 
+// SimpleExpectation holds an arbitrary check func, and implements the Expectation interface.
+type SimpleExpectation struct {
+	check       func(State) Verdict
+	description string
+}
+
+// Check invokes e.check.
+func (e SimpleExpectation) Check(s State) Verdict {
+	return e.check(s)
+}
+
+// Description returns e.descriptin.
+func (e SimpleExpectation) Description() string {
+	return e.description
+}
+
+// NoOutstandingWork asserts that there is no work initiated using the LSP
+// $/progress API that has not completed.
+func NoOutstandingWork() SimpleExpectation {
+	check := func(s State) Verdict {
+		if len(s.outstandingWork) == 0 {
+			return Met
+		}
+		return Unmet
+	}
+	return SimpleExpectation{
+		check:       check,
+		description: "no outstanding work",
+	}
+}
+
 // LogExpectation is an expectation on the log messages received by the editor
 // from gopls.
 type LogExpectation struct {