Fix flaky TestPartialStateJoin/Outgoing_device_list_updates from a leave-handler race

`WithWaitForLeave` registered a PDU handler (whose callback feeds
`leaveChannel`) but then decided whether to wait by peeking at
`room.CurrentState`. Because an incoming PDU is added to the room
(`HandleTransactionRequests` calls `room.AddEvent`, updating
`CurrentState`) *before* the PDU callback runs, the test goroutine could
observe the updated state, take the "already seen" shortcut and return,
firing the deferred `removePDUHandler()` in the window before the
callback ran. The now-handler-less callback then flagged the (expected)
leave as an unexpected PDU, failing the test.

Wait on the channel fed by the handler instead of polling
`room.CurrentState`, so the handler stays registered until its own
callback has run.

Every caller passes a joined user and an action that performs the leave, so the
leave always arrives fresh via that callback. Make this contract
explicit: snapshot membership before the action and fail if the user has
already left (rather than silently treating it as "already seen"), and
fail on timeout, since neither should happen in correct usage. The
action still runs in the already-left case so any cleanup it performs is
preserved.

The `HandleTransactionRequests` ordering (callback after `room.AddEvent`) is left
unchanged, as other tests rely on it (e.g.
`federation_redaction_test.go` reads `serverRoom.Timeline` immediately
after a callback-driven waiter).

Author: erikjohnston

tests/msc3902/federation_room_join_partial_state_test.go

@@ -141,23 +141,34 @@ func (s *server) WithWaitForLeave(
 	)
 	defer removePDUHandler()
 
-	leaveAction()
-
+	// `WithWaitForLeave` expects the user to be in the room and to leave as a
+	// result of the action, so that the leave arrives as a PDU our handler
+	// matches. Snapshot the membership *before* the action: if the user has
+	// already left, no such PDU is coming, which is a test bug rather than
+	// something to wait for.
 	memberEvent := room.CurrentState("m.room.member", userID)
-	membership := ""
+	membership := "leave"
 	if memberEvent != nil {
 		membership, _ = memberEvent.Membership()
 	}
-	if membership == "leave" {
-		t.Logf("%s has already seen %s leave test room %s.", s.ServerName(), userID, room.RoomID)
-	} else {
-		select {
-		case <-leaveChannel:
-			t.Logf("%s saw %s leave test room %s.", s.ServerName(), userID, room.RoomID)
-			break
-		case <-time.After(1 * time.Second):
-			t.Errorf("%s timed out waiting for %s to leave test room %s.", s.ServerName(), userID, room.RoomID)
-		}
+	alreadyLeft := membership == "leave"
+	if alreadyLeft {
+		t.Errorf("%s: %s had already left test room %s before WithWaitForLeave ran.", s.ServerName(), userID, room.RoomID)
+	}
+
+	leaveAction()
+
+	// Wait on the channel fed by our PDU handler rather than polling
+	// `room.CurrentState`: the room's current state is updated (by
+	// `room.AddEvent`) *before* the PDU callback runs, so returning on a
+	// `CurrentState` check could deregister our handler in the window before
+	// the callback fires, making the (expected) leave look unexpected to
+	// `HandleTransactionRequests`.
+	select {
+	case <-leaveChannel:
+		t.Logf("%s saw %s leave test room %s.", s.ServerName(), userID, room.RoomID)
+	case <-time.After(1 * time.Second):
+		t.Errorf("%s timed out waiting for %s to leave test room %s.", s.ServerName(), userID, room.RoomID)
 	}
 }