Fix tests

2 files changed, 3 insertions, 2 deletions

diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py
index 8331f6422e..75f259ee4e 100644
--- a/synapse/handlers/user_directory.py
+++ b/synapse/handlers/user_directory.py
@@ -130,7 +130,7 @@ class UserDirectoyHandler(object):
     def _handle_intial_room(self, room_id):
         """Called when we initially fill out user_directory one room at a time
         """
-        is_in_room = yield self.store.get_is_host_in_room(room_id, self.server_name)
+        is_in_room = yield self.state.get_is_host_in_room(room_id, self.server_name)
         if not is_in_room:
             return
 
@@ -232,7 +232,7 @@ class UserDirectoyHandler(object):
                 if not change:
                     # Need to check if the server left the room entirely, if so
                     # we might need to remove all the users in that room
-                    is_in_room = yield self.store.get_is_host_in_room(
+                    is_in_room = yield self.state.get_is_host_in_room(
                         room_id, self.server_name,
                     )
                     if not is_in_room:
diff --git a/synapse/storage/user_directory.py b/synapse/storage/user_directory.py
index 9137fc24ea..348064436f 100644
--- a/synapse/storage/user_directory.py
+++ b/synapse/storage/user_directory.py
@@ -221,6 +221,7 @@ class UserDirectoryStore(SQLBaseStore):
             txn.execute(sql, (prev_stream_id,))
 
             total = 0
+            max_stream_id = prev_stream_id
             for max_stream_id, count in txn:
                 total += count
                 if total > 50: