Restructure update membership

src/github.com/matrix-org/dendrite/roomserver/input/membership.go

@@ -97,24 +97,43 @@ func updateMembership(
 		return nil, err
 	}
 
-	return updateWithNewMembership(mu, new, add, updates)
+	switch new {
+	case "invite":
+		return updateToInviteMembership(mu, add, updates)
+	case "join":
+		return updateToJoinMembership(mu, add, updates)
+	case "leave", "ban":
+		return updateToLeaveMembership(mu, add, new, updates)
+	default:
+		panic(fmt.Errorf(
+			"input: membership %q is not one of the allowed values", new,
+		))
+	}
 }
 
-func updateWithNewMembership(
+func updateToInviteMembership(
-	mu types.MembershipUpdater, newMembership string, add *gomatrixserverlib.Event,
+	mu types.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
-	updates []api.OutputEvent,
 ) ([]api.OutputEvent, error) {
-	switch newMembership {
-	case "invite":
 	needsSending, err := mu.SetToInvite(*add)
 	if err != nil {
 		return nil, err
 	}
 	if needsSending {
-			updates = appendInviteUpdate(updates, add)
+		onie := api.OutputNewInviteEvent{
+			Event: *add,
+		}
+		updates = append(updates, api.OutputEvent{
+			Type:           api.OutputTypeNewInviteEvent,
+			NewInviteEvent: &onie,
+		})
 	}
 	return updates, nil
-	case "join":
+}
+
+func updateToJoinMembership(
+	mu types.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
+) ([]api.OutputEvent, error) {
+
 	if mu.IsJoin() {
 		return updates, nil
 	}
@@ -122,42 +141,10 @@ func updateWithNewMembership(
 	if err != nil {
 		return nil, err
 	}
-		return appendRetireUpdates(updates, add, newMembership, retired), nil
-	case "leave", "ban":
-		if mu.IsLeave() {
-			return updates, nil
-		}
-		retired, err := mu.SetToLeave(add.Sender())
-		if err != nil {
-			mu.SetToLeave(add.Sender())
-		}
-		return appendRetireUpdates(updates, add, newMembership, retired), nil
-	default:
-		panic(fmt.Errorf(
-			"input: membership %q is not one of the allowed values", newMembership,
-		))
-	}
-}
-
-func appendInviteUpdate(
-	updates []api.OutputEvent, add *gomatrixserverlib.Event,
-) []api.OutputEvent {
-	onie := api.OutputNewInviteEvent{
-		Event: *add,
-	}
-	return append(updates, api.OutputEvent{
-		Type:           api.OutputTypeNewInviteEvent,
-		NewInviteEvent: &onie,
-	})
-}
-
-func appendRetireUpdates(
-	updates []api.OutputEvent, add *gomatrixserverlib.Event, membership string, retired []string,
-) []api.OutputEvent {
 	for _, eventID := range retired {
 		orie := api.OutputRetireInviteEvent{
 			EventID:    eventID,
-			Membership: membership,
+			Membership: "join",
 		}
 		if add != nil {
 			orie.RetiredByEventID = add.EventID()
@@ -167,7 +154,35 @@ func appendRetireUpdates(
 			RetireInviteEvent: &orie,
 		})
 	}
-	return updates
+	return updates, nil
+}
+
+func updateToLeaveMembership(
+	mu types.MembershipUpdater, add *gomatrixserverlib.Event,
+	newMembership string, updates []api.OutputEvent,
+) ([]api.OutputEvent, error) {
+
+	if mu.IsLeave() {
+		return updates, nil
+	}
+	retired, err := mu.SetToLeave(add.Sender())
+	if err != nil {
+		return nil, err
+	}
+	for _, eventID := range retired {
+		orie := api.OutputRetireInviteEvent{
+			EventID:    eventID,
+			Membership: newMembership,
+		}
+		if add != nil {
+			orie.RetiredByEventID = add.EventID()
+		}
+		updates = append(updates, api.OutputEvent{
+			Type:              api.OutputTypeRetireInviteEvent,
+			RetireInviteEvent: &orie,
+		})
+	}
+	return updates, nil
 }
 
 type stateChange struct {