Add ServerRoom.MustHaveMembershipForUser (#229)

* Add ServerRoom.MustHaveMembershipForUser

As per https://github.com/matrix-org/complement/pull/131/files#r658674594

* Docstring

Author: kegsay

internal/federation/server_room.go

@@ -3,6 +3,7 @@ package federation
 import (
 	"encoding/json"
 	"fmt"
+	"testing"
 
 	"github.com/matrix-org/gomatrixserverlib"
 
@@ -112,6 +113,22 @@ func (r *ServerRoom) AuthChain() (chain []*gomatrixserverlib.Event) {
 	return
 }
 
+// Check that the user currently has the membership provided in this room. Fails the test if not.
+func (r *ServerRoom) MustHaveMembershipForUser(t *testing.T, userID, wantMembership string) {
+	t.Helper()
+	state := r.CurrentState("m.room.member", userID)
+	if state == nil {
+		t.Fatalf("no membership state for %s", userID)
+	}
+	m, err := state.Membership()
+	if err != nil {
+		t.Fatalf("m.room.member event exists for %s but cannot read membership field: %s", userID, err)
+	}
+	if m != wantMembership {
+		t.Fatalf("incorrect membership state for %s: got %s, want %s", userID, m, wantMembership)
+	}
+}
+
 func initialPowerLevelsContent(roomCreator string) (c gomatrixserverlib.PowerLevelContent) {
 	c.Defaults()
 	c.Events = map[string]int64{

tests/federation_room_invite_test.go

@@ -1,7 +1,6 @@
 package tests
 
 import (
-	"fmt"
 	"testing"
 	"time"
 
@@ -52,32 +51,11 @@ func TestFederationRejectInvite(t *testing.T) {
 	waiter = NewWaiter()
 	alice.InviteRoom(t, roomID, charlie.UserID)
 	waiter.Wait(t, 5*time.Second)
-	if err := checkMembershipForUser(room, charlie.UserID, "invite"); err != nil {
-		t.Errorf("Membership state for charlie after invite: %v", err)
-	}
+	room.MustHaveMembershipForUser(t, charlie.UserID, "invite")
 
 	// Charlie rejects the invite; Delia should see the rejection.
 	waiter = NewWaiter()
 	charlie.LeaveRoom(t, roomID)
 	waiter.Wait(t, 5*time.Second)
-	if err := checkMembershipForUser(room, charlie.UserID, "leave"); err != nil {
-		t.Errorf("Membership state for charlie after reject: %v", err)
-	}
-}
-
-func checkMembershipForUser(room *federation.ServerRoom, userID, wantMembership string) (err error) {
-	state := room.CurrentState("m.room.member", userID)
-	if state == nil {
-		err = fmt.Errorf("no membership state for %s", userID)
-		return
-	}
-	m, err := state.Membership()
-	if err != nil {
-		return
-	}
-	if m != wantMembership {
-		err = fmt.Errorf("incorrect membership state: got %s, want %s", m, wantMembership)
-		return
-	}
-	return
+	room.MustHaveMembershipForUser(t, charlie.UserID, "leave")
 }