Faster joins: Refactor handling of servers in room (#14954)

Ensure that the list of servers in a partial state room always contains
the server we joined off.

Also refactor `get_partial_state_servers_at_join` to return `None` when
the given room is no longer partial stated, to explicitly indicate when
the room has partial state. Otherwise it's not clear whether an empty
list means that the room has full state, or the room is partial stated,
but the server we joined off told us that there are no servers in the
room.

Signed-off-by: Sean Quah <seanq@matrix.org>

2 files changed, 2 insertions, 2 deletions

diff --git a/tests/handlers/test_federation.py b/tests/handlers/test_federation.py
index c1558c40c3..57675fa407 100644
--- a/tests/handlers/test_federation.py
+++ b/tests/handlers/test_federation.py
@@ -656,7 +656,7 @@ class PartialJoinTestCase(unittest.FederatingHomeserverTestCase):
                         EVENT_INVITATION_MEMBERSHIP,
                     ],
                     partial_state=True,
-                    servers_in_room=["example.com"],
+                    servers_in_room={"example.com"},
                 )
             )
         )
diff --git a/tests/handlers/test_room_member.py b/tests/handlers/test_room_member.py
index 6bbfd5dc84..6a38893b68 100644
--- a/tests/handlers/test_room_member.py
+++ b/tests/handlers/test_room_member.py
@@ -171,7 +171,7 @@ class TestJoinsLimitedByPerRoomRateLimiter(FederatingHomeserverTestCase):
                     state=[create_event],
                     auth_chain=[create_event],
                     partial_state=False,
-                    servers_in_room=[],
+                    servers_in_room=frozenset(),
                 )
             )
         )