Fix some of the tests to reflect changes in the storage layer.

3 files changed, 11 insertions, 39 deletions

diff --git a/tests/handlers/test_room.py b/tests/handlers/test_room.py
index fd2d66db38..bfdde6135b 100644
--- a/tests/handlers/test_room.py
+++ b/tests/handlers/test_room.py
@@ -40,7 +40,7 @@ class RoomMemberHandlerTestCase(unittest.TestCase):
             self.hostname,
             db_pool=None,
             datastore=NonCallableMock(spec_set=[
-                "store_room_member",
+                "persist_event",
                 "get_joined_hosts_for_room",
                 "get_room_member",
                 "get_room",
@@ -97,7 +97,7 @@ class RoomMemberHandlerTestCase(unittest.TestCase):
         )
 
         store_id = "store_id_fooo"
-        self.datastore.store_room_member.return_value = defer.succeed(store_id)
+        self.datastore.persist_event.return_value = defer.succeed(store_id)
 
         # Actual invocation
         yield self.room_member_handler.change_membership(event)
@@ -110,12 +110,8 @@ class RoomMemberHandlerTestCase(unittest.TestCase):
             set(event.destinations)
         )
 
-        self.datastore.store_room_member.assert_called_once_with(
-            user_id=target_user_id,
-            sender=user_id,
-            room_id=room_id,
-            content=content,
-            membership=Membership.INVITE,
+        self.datastore.persist_event.assert_called_once_with(
+            event
         )
         self.notifier.on_new_room_event.assert_called_once_with(
                 event, store_id)
@@ -149,7 +145,7 @@ class RoomMemberHandlerTestCase(unittest.TestCase):
         )
 
         store_id = "store_id_fooo"
-        self.datastore.store_room_member.return_value = defer.succeed(store_id)
+        self.datastore.persist_event.return_value = defer.succeed(store_id)
         self.datastore.get_room.return_value = defer.succeed(1)  # Not None.
 
         prev_state = NonCallableMock()
@@ -171,12 +167,8 @@ class RoomMemberHandlerTestCase(unittest.TestCase):
             set(event.destinations)
         )
 
-        self.datastore.store_room_member.assert_called_once_with(
-            user_id=target_user_id,
-            sender=user_id,
-            room_id=room_id,
-            content=content,
-            membership=Membership.JOIN,
+        self.datastore.persist_event.assert_called_once_with(
+            event
         )
         self.notifier.on_new_room_event.assert_called_once_with(
                 event, store_id)
diff --git a/tests/test_state.py b/tests/test_state.py
index aaf873a856..e64d15a3a2 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -36,7 +36,7 @@ class StateTestCase(unittest.TestCase):
             "get_unresolved_state_tree",
             "update_current_state",
             "get_latest_pdus_in_context",
-            "get_current_state",
+            "get_current_state_pdu",
             "get_pdu",
         ])
         self.replication = Mock(spec=["get_pdu"])
@@ -247,7 +247,7 @@ class StateTestCase(unittest.TestCase):
         pdus = [tup]
 
         self.persistence.get_latest_pdus_in_context.return_value = pdus
-        self.persistence.get_current_state.return_value = state_pdu
+        self.persistence.get_current_state_pdu.return_value = state_pdu
 
         yield self.state.handle_new_event(event)
 
diff --git a/tests/utils.py b/tests/utils.py
index 20a63316fd..717b81e6ee 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -221,34 +221,14 @@ class MemoryDataStore(object):
                     content=content)
         self.paths_to_content[path] = data
 
-    def get_message_stream(self, user_id=None, from_key=None, to_key=None,
+    def get_room_events_stream(self, user_id=None, from_key=None, to_key=None,
                             room_id=None, limit=0, with_feedback=False):
         return ([], from_key)  # TODO
 
-    def get_room_member_stream(self, user_id=None, from_key=None, to_key=None):
-        return ([], from_key)  # TODO
-
-    def get_feedback_stream(self, user_id=None, from_key=None, to_key=None,
-                            room_id=None, limit=0):
-        return ([], from_key)  # TODO
-
-    def get_room_data_stream(self, user_id=None, from_key=None, to_key=None,
-                            room_id=None, limit=0):
-        return ([], from_key)  # TODO
-
     def to_events(self, data_store_list):
         return data_store_list  # TODO
 
-    def get_max_message_id(self):
-        return 0  # TODO
-
-    def get_max_feedback_id(self):
-        return 0  # TODO
-
-    def get_max_room_member_id(self):
-        return 0  # TODO
-
-    def get_max_room_data_id(self):
+    def get_room_events_max_id(self):
         return 0  # TODO
 
     def get_joined_hosts_for_room(self, room_id):