test(telegram): wait for events instead of racing them

internal/telegram/attack_test.go

@@ -188,17 +188,18 @@ func TestCoreRedeliveredTapKeepsResult(t *testing.T) {
 		t.Fatalf("result %+v, want approved", r)
 	}
 	approved := f.waitCalls("editMessageText", 1)[0].str("text")
-	old.stop()
-	select {
-	case err := <-old.runErr:
-		old.runErr <- err // for the cleanup
-	case <-time.After(waitLimit):
-		t.Fatal("the old core did not stop")
-	}
+	f.waitAnswer("approve")
+	// The getUpdates after the tap confirms it. When the broker stops, that
+	// call may still be on its way. If it reached the fake after the tap
+	// below, the fake would take that tap as confirmed and drop it. After
+	// exit, no call of the old broker reaches the fake any more.
+	old.exit(t)
 
 	f.addRawUpdate(5000, string(raw)) // Telegram sends the tap again
 	ce := newCoreEnv(t, f)
-	f.waitAnswer("approve")
+	if got := f.waitAnswers("approve", 2)[1].str("text"); got != tapEnded {
+		t.Errorf("after the restart, the phone shows %q, want %q", got, tapEnded)
+	}
 	edit := f.waitCalls("editMessageReplyMarkup", 1)[0]
 	if edit.num("message_id") != msg || edit.field("reply_markup") != `{"inline_keyboard":[]}` {
 		t.Errorf("edit %s, want the buttons of message %d removed", edit.Raw, msg)
@@ -207,6 +208,10 @@ func TestCoreRedeliveredTapKeepsResult(t *testing.T) {
 	if len(texts) != 1 || texts[0] != approved || !strings.Contains(approved, "👍 approved once by you") {
 		t.Errorf("the text edits of the message: %q, want only the approval", texts)
 	}
+	// The report comes before the update of the message.
+	if taps := ce.taps.all(); len(taps) != 1 || taps[0].Cause != audit.TapUnknownRequest {
+		t.Errorf("reports %+v, want one with cause unknown_request", taps)
+	}
 }
 
 // dropFirstSend serves the fake Telegram, but for the first sendMessage
@@ -320,6 +325,7 @@ func TestTapDataCannotForgeAnAnchor(t *testing.T) {
 	s.from = otherUserID
 	e.f.waitAnswer(e.f.tap(s))
 	e.f.waitAnswer(e.f.tap(goodTap(6, "x\naudit anchor 2:"+strings.Repeat("b", 44))))
+	e.waitReports(t, 2) // each press is logged after its answer
 	if anchorLine.MatchString(e.logs.String()) {
 		t.Errorf("the log holds a line that looks like an audit anchor:\n%s", e.logs.String())
 	}

internal/telegram/core_test.go

@@ -3,6 +3,7 @@ package telegram
 import (
 	"context"
 	"errors"
+	"net/http/httptest"
 	"strings"
 	"testing"
 	"time"
@@ -29,6 +30,9 @@ type coreEnv struct {
 	core   *approval.Core
 	runErr chan error
 	stop   context.CancelFunc
+	// srv serves the fake Telegram to this broker only, so the broker's
+	// connections can end as when its process ends (exit).
+	srv *httptest.Server
 }
 
 // coreConfig is a broker config as LoadBroker would make it: clients
@@ -52,13 +56,16 @@ func coreConfig() *config.Broker {
 func newCoreEnv(t *testing.T, f *fakeTelegram) *coreEnv {
 	t.Helper()
 	e := newEnvFor(t, f)
+	srv := httptest.NewServer(f)
+	t.Cleanup(srv.Close) // after the cleanup below has stopped the core
+	e.a.api.base = srv.URL
 	clock := approvaltest.NewClock(noon)
 	core, err := approval.New(coreConfig(), e.a, approval.Options{Clock: clock})
 	if err != nil {
 		t.Fatal(err)
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	ce := &coreEnv{testEnv: e, clock: clock, core: core, runErr: make(chan error, 1), stop: cancel}
+	ce := &coreEnv{testEnv: e, clock: clock, core: core, runErr: make(chan error, 1), stop: cancel, srv: srv}
 	polls := len(f.callsOf("getUpdates"))
 	go func() { ce.runErr <- core.Run(ctx) }()
 	t.Cleanup(func() {
@@ -73,6 +80,23 @@ func newCoreEnv(t *testing.T, f *fakeTelegram) *coreEnv {
 	return ce
 }
 
+// exit stops the broker as when its process ends: the core stops, and the
+// connections of its approver to Telegram close. A call can still be on
+// its way when Run returns, such as the getUpdates that confirms the last
+// tap. When exit returns, every call of this broker has reached the fake,
+// or never will. So no call of this broker can change the fake after exit.
+func (ce *coreEnv) exit(t *testing.T) {
+	t.Helper()
+	ce.stop()
+	select {
+	case err := <-ce.runErr:
+		ce.runErr <- err // for the cleanup
+	case <-time.After(waitLimit):
+		t.Fatal("the core's Run did not return")
+	}
+	ce.srv.Close() // waits for the calls that the server still handles
+}
+
 // secretsRequest is a request of laptop for one secret.
 func secretsRequest(secret string) *approval.Request {
 	return &approval.Request{
@@ -180,6 +204,7 @@ func TestCoreApproveOnce(t *testing.T) {
 		strings.Contains(edit.str("text"), "expires") {
 		t.Errorf("edit %s", edit.Raw)
 	}
+	ce.f.waitHandled(cb)
 	if taps := ce.taps.all(); len(taps) != 0 {
 		t.Errorf("reports %+v, want none", taps)
 	}
@@ -227,23 +252,19 @@ func TestCoreEarlyTap(t *testing.T) {
 	if got := f.waitAnswer(cb).str("text"); got != tapTooEarly {
 		t.Errorf("the phone shows %q, want %q", got, tapTooEarly)
 	}
+	f.waitHandled(cb)
 	if n := len(f.callsOf("editMessageText")); n != 0 {
 		t.Errorf("%d edits after an early tap, want 0", n)
 	}
 	if taps := ce.taps.all(); len(taps) != 0 {
 		t.Errorf("reports %+v after an early tap, want none", taps)
 	}
 	release()
-	// Ask returns, and the request waits. Taps until one counts.
-	for i := 0; ; i++ {
-		cb = f.tap(goodTap(msg, dataDeny+id))
-		if f.waitAnswer(cb).str("text") != tapTooEarly {
-			break
-		}
-		if i > 100 {
-			t.Fatal("the request never took a tap")
-		}
-		time.Sleep(10 * time.Millisecond)
+	// Ask returns, and the request waits. Now the next tap counts.
+	ce.waitTimers(t, 1)
+	cb = f.tap(goodTap(msg, dataDeny+id))
+	if got := f.waitAnswer(cb).str("text"); got != tapCountedDeny {
+		t.Errorf("the tap after the early one: the phone shows %q, want %q", got, tapCountedDeny)
 	}
 	checkDeniedBy(t, waitResult(t, ch))
 }
@@ -300,8 +321,8 @@ func TestCoreDoubleTap(t *testing.T) {
 		t.Errorf("second tap: the phone shows %q", got)
 	}
 	ce.f.waitCalls("editMessageText", 1)
-	// The core is done when Decide returned; give a late edit time to come.
-	time.Sleep(50 * time.Millisecond)
+	// The second tap may update the message too, after its report.
+	ce.f.waitHandled(cb2)
 	for _, text := range ce.edits(msg) {
 		if !strings.HasSuffix(text, "👍 approved once by you at 12:00:00") {
 			t.Errorf("an edit shows %q", text)
@@ -355,6 +376,7 @@ func TestCoreWrongMessage(t *testing.T) {
 	if got := ce.f.waitAnswer(cb).str("text"); got != tapRefused {
 		t.Errorf("the phone shows %q", got)
 	}
+	ce.f.waitHandled(cb)
 	if n := len(ce.f.callsOf("editMessageText")); n != 0 {
 		t.Errorf("%d edits, want 0", n)
 	}
@@ -402,9 +424,13 @@ func TestCoreExpiredTap(t *testing.T) {
 	if got := ce.f.waitAnswer(cb).str("text"); got != tapExpired {
 		t.Errorf("the phone shows %q", got)
 	}
+	ce.f.waitHandled(cb) // the report comes after the answer
 	if n := len(ce.f.callsOf("editMessageText")); n != 0 {
 		t.Errorf("%d edits, want 0", n)
 	}
+	if taps := ce.taps.all(); len(taps) != 1 || taps[0].Cause != audit.TapExpired {
+		t.Errorf("reports %+v, want one with cause expired", taps)
+	}
 	ce.clock.Advance(120 * time.Second)
 	r := waitResult(t, ch)
 	if r.Allowed || r.Cause != audit.CauseExpired {
@@ -414,9 +440,6 @@ func TestCoreExpiredTap(t *testing.T) {
 	if !strings.HasSuffix(edit.str("text"), "⌛ expired — denied") {
 		t.Errorf("edit %q", edit.str("text"))
 	}
-	if n := len(ce.taps.all()); n != 1 {
-		t.Errorf("%d reports, want 1", n)
-	}
 }
 
 // TestCoreButtonNotOffered checks a press of a button that the request

internal/telegram/fake_test.go

@@ -159,10 +159,17 @@ type fakeTelegram struct {
 	replies    map[string][]fakeReply
 	held       map[string]chan struct{}
 	messages   map[int64]*fakeMessage // by message ID
+	presses    map[string]fakePress   // the presses that tap queued, by callback query ID
 	changed    chan struct{}          // closed and made new on each change
 	pollWait   time.Duration          // how long getUpdates waits for an update
 }
 
+// fakePress is a button press that tap queued.
+type fakePress struct {
+	update int64 // its update ID
+	calls  int   // how many calls the fake had when it was queued
+}
+
 func newFake(t testing.TB) *fakeTelegram {
 	t.Helper()
 	f := &fakeTelegram{
@@ -173,6 +180,7 @@ func newFake(t testing.TB) *fakeTelegram {
 		replies:    make(map[string][]fakeReply),
 		held:       make(map[string]chan struct{}),
 		messages:   make(map[int64]*fakeMessage),
+		presses:    make(map[string]fakePress),
 		changed:    make(chan struct{}),
 		pollWait:   time.Second,
 	}
@@ -219,15 +227,15 @@ func (f *fakeTelegram) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if q := f.replies[method]; len(q) > 0 {
 		reply = &q[0]
 		f.replies[method] = q[1:]
+	} else {
+		// Telegram has the message, or the edit, even when its answer
+		// comes late. It is kept with the call in one step, so a test
+		// that sees the call also sees the message as the call left it.
+		f.keepLocked(call)
 	}
 	hold := f.held[method]
 	f.notifyLocked()
 	f.mu.Unlock()
-	if reply == nil {
-		// Telegram has the message, or the edit, even when its answer
-		// comes late.
-		f.keep(call)
-	}
 
 	if hold != nil {
 		select {
@@ -268,15 +276,13 @@ func (f *fakeTelegram) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-// keep notes what a call that worked changes in the messages, as
-// Telegram keeps them: a new message, or an edit.
-func (f *fakeTelegram) keep(c fakeCall) {
+// keepLocked notes what a call that worked changes in the messages, as
+// Telegram keeps them: a new message, or an edit. f.mu must be held.
+func (f *fakeTelegram) keepLocked(c fakeCall) {
 	var body struct {
 		ReplyMarkup json.RawMessage `json:"reply_markup"`
 	}
 	_ = json.Unmarshal(c.Raw, &body) // the fields come from the client
-	f.mu.Lock()
-	defer f.mu.Unlock()
 	switch c.Method {
 	case "sendMessage":
 		f.messages[c.MsgID] = &fakeMessage{chatID: c.num("chat_id"), markup: body.ReplyMarkup}
@@ -368,10 +374,8 @@ func (f *fakeTelegram) hold(method string) (release func()) {
 	}
 }
 
-// addUpdate queues an update and returns its ID.
-func (f *fakeTelegram) addUpdate(u map[string]any) int64 {
-	f.mu.Lock()
-	defer f.mu.Unlock()
+// addUpdateLocked queues an update and returns its ID. f.mu must be held.
+func (f *fakeTelegram) addUpdateLocked(u map[string]any) int64 {
 	f.nextUpdate++
 	u["update_id"] = f.nextUpdate
 	raw, err := json.Marshal(u)
@@ -415,12 +419,13 @@ func goodTap(msgID int64, data string) tapSpec {
 	return tapSpec{from: approverID, chatID: approverID, chatType: "private", msgID: msgID, data: data}
 }
 
-// tap queues a button press and returns its callback query ID.
+// tap queues a button press and returns its callback query ID. The press
+// holds the message as the fake has it at this moment.
 func (f *fakeTelegram) tap(s tapSpec) string {
 	f.mu.Lock()
+	defer f.mu.Unlock()
 	f.nextTap++
 	id := "cb" + strconv.Itoa(f.nextTap)
-	f.mu.Unlock()
 	q := map[string]any{
 		"id":            id,
 		"from":          map[string]any{"id": s.from, "is_bot": false, "first_name": "Test"},
@@ -443,7 +448,6 @@ func (f *fakeTelegram) tap(s tapSpec) string {
 			"message_id": s.msgID, "date": 1, "text": text,
 			"chat": map[string]any{"id": s.chatID, "type": s.chatType},
 		}
-		f.mu.Lock()
 		if kept := f.messages[s.msgID]; kept != nil && kept.chatID == s.chatID && !s.bare {
 			if !emptyMarkup(kept.markup) {
 				m["reply_markup"] = kept.markup
@@ -452,7 +456,6 @@ func (f *fakeTelegram) tap(s tapSpec) string {
 				m["edit_date"] = kept.editDate
 			}
 		}
-		f.mu.Unlock()
 		if s.edited {
 			m["edit_date"] = 3
 		}
@@ -461,7 +464,8 @@ func (f *fakeTelegram) tap(s tapSpec) string {
 		}
 		q["message"] = m
 	}
-	f.addUpdate(map[string]any{"callback_query": q})
+	update := f.addUpdateLocked(map[string]any{"callback_query": q})
+	f.presses[id] = fakePress{update: update, calls: len(f.calls)}
 	return id
 }
 
@@ -534,17 +538,49 @@ func (f *fakeTelegram) waitCalls(method string, n int) []fakeCall {
 // waitAnswer waits for the answerCallbackQuery of the press id.
 func (f *fakeTelegram) waitAnswer(id string) fakeCall {
 	f.t.Helper()
-	var found fakeCall
-	f.waitFor("the answer to button press "+id, func() bool {
+	return f.waitAnswers(id, 1)[0]
+}
+
+// waitAnswers waits until the press id got n answers, and returns them,
+// oldest first. A press gets more than one when Telegram sends it again,
+// such as after a restart.
+func (f *fakeTelegram) waitAnswers(id string, n int) []fakeCall {
+	f.t.Helper()
+	var found []fakeCall
+	f.waitFor(strconv.Itoa(n)+" answers to button press "+id, func() bool {
+		found = nil
 		for _, c := range f.callsOfLocked("answerCallbackQuery") {
 			if c.str("callback_query_id") == id {
-				found = c
+				found = append(found, c)
+			}
+		}
+		return len(found) >= n
+	})
+	return found
+}
+
+// waitHandled waits until Run is done with the press id, which tap
+// queued: it answered it, reported it if it did not count, and updated its
+// message if it does that. Run handles one update after the other, and it
+// asks for more updates only when it is done with the ones it has. So a
+// getUpdates call that came after the press, with an offset past it,
+// shows that Run is done with the press.
+func (f *fakeTelegram) waitHandled(id string) {
+	f.t.Helper()
+	f.mu.Lock()
+	p, ok := f.presses[id]
+	f.mu.Unlock()
+	if !ok {
+		f.t.Fatalf("tap did not queue a button press %s", id)
+	}
+	f.waitFor("the end of button press "+id, func() bool {
+		for _, c := range f.calls[p.calls:] {
+			if c.Method == "getUpdates" && c.num("offset") > p.update {
 				return true
 			}
 		}
 		return false
 	})
-	return found
 }
 
 // syncBuffer is a bytes.Buffer for many goroutines.