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,21 +141,42 @@ func (s *server) WithWaitForLeave(
 	)
 	defer removePDUHandler()
 
+	currentMembership := func() string {
+		memberEvent := room.CurrentState("m.room.member", userID)
+		if memberEvent == nil {
+			return ""
+		}
+		membership, _ := memberEvent.Membership()
+		return membership
+	}
+
+	// Snapshot the membership *before* running the action. If the user has
+	// already left, the leave PDU predates our handler, so no PDU callback is
+	// pending for it and there is nothing to wait for.
+	alreadyLeft := currentMembership() == "leave"
+
 	leaveAction()
 
-	memberEvent := room.CurrentState("m.room.member", userID)
-	membership := ""
-	if memberEvent != nil {
-		membership, _ = memberEvent.Membership()
-	}
-	if membership == "leave" {
+	if alreadyLeft {
 		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):
+		return
+	}
+
+	// Otherwise the leave is expected to arrive as a result of the action. 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):
+		// Fall back to the room state in case the leave was applied via a path
+		// that doesn't fire the PDU callback (e.g. added to the room directly).
+		if currentMembership() == "leave" {
+			t.Logf("%s has already seen %s leave test room %s.", s.ServerName(), userID, room.RoomID)
+		} else {
 			t.Errorf("%s timed out waiting for %s to leave test room %s.", s.ServerName(), userID, room.RoomID)
 		}
 	}