2 files changed, 9 insertions, 8 deletions
diff --git a/tests/rest/test_rooms.py b/tests/rest/test_rooms.py
index b432cf254e..cdaf948a3b 100644
--- a/tests/rest/test_rooms.py
+++ b/tests/rest/test_rooms.py
@@ -51,7 +51,7 @@ class RoomPermissionsTestCase(RestTestCase):
         persistence_service.get_latest_pdus_in_context.return_value = []
 
         hs = HomeServer(
-            "test",
+            "red",
             db_pool=None,
             http_client=None,
             datastore=MemoryDataStore(),
@@ -398,7 +398,7 @@ class RoomsMemberListTestCase(RestTestCase):
         persistence_service.get_latest_pdus_in_context.return_value = []
 
         hs = HomeServer(
-            "test",
+            "red",
             db_pool=None,
             http_client=None,
             datastore=MemoryDataStore(),
@@ -476,7 +476,7 @@ class RoomsCreateTestCase(RestTestCase):
         persistence_service.get_latest_pdus_in_context.return_value = []
 
         hs = HomeServer(
-            "test",
+            "red",
             db_pool=None,
             http_client=None,
             datastore=MemoryDataStore(),
@@ -566,7 +566,7 @@ class RoomTopicTestCase(RestTestCase):
         persistence_service.get_latest_pdus_in_context.return_value = []
 
         hs = HomeServer(
-            "test",
+            "red",
             db_pool=None,
             http_client=None,
             datastore=MemoryDataStore(),
@@ -669,7 +669,7 @@ class RoomMemberStateTestCase(RestTestCase):
         persistence_service.get_latest_pdus_in_context.return_value = []
 
         hs = HomeServer(
-            "test",
+            "red",
             db_pool=None,
             http_client=None,
             datastore=MemoryDataStore(),
@@ -794,7 +794,7 @@ class RoomMessagesTestCase(RestTestCase):
         persistence_service.get_latest_pdus_in_context.return_value = []
 
         hs = HomeServer(
-            "test",
+            "red",
             db_pool=None,
             http_client=None,
             datastore=MemoryDataStore(),
diff --git a/tests/utils.py b/tests/utils.py
index 37b759febc..ea7d6893c6 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -188,8 +188,9 @@ class MemoryDataStore(object):
 
     def get_rooms_for_user_where_membership_is(self, user_id, membership_list):
         return [
-            r for r in self.members
-            if self.members[r].get(user_id).membership in membership_list
+            self.members[r].get(user_id) for r in self.members
+            if user_id in self.members[r] and
+                self.members[r][user_id].membership in membership_list
         ]
 
     def get_room_events_stream(self, user_id=None, from_key=None, to_key=None,