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
`CurrentState`, so the handler stays registered until its own callback
has run. Snapshot membership before the action to distinguish a
pre-existing leave (nothing to wait for) from a fresh one, and keep a
`CurrentState` check only as a timeout fallback for leaves applied via
non-callback paths.

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)
 		}
 	}