Various clean ups to room stream tokens. (#8423)

2 files changed, 8 insertions, 10 deletions

diff --git a/tests/rest/client/v1/test_rooms.py b/tests/rest/client/v1/test_rooms.py
index 0a567b032f..a3287011e9 100644
--- a/tests/rest/client/v1/test_rooms.py
+++ b/tests/rest/client/v1/test_rooms.py
@@ -902,15 +902,15 @@ class RoomMessageListTestCase(RoomBase):
 
         # Send a first message in the room, which will be removed by the purge.
         first_event_id = self.helper.send(self.room_id, "message 1")["event_id"]
-        first_token = self.get_success(
-            store.get_topological_token_for_event(first_event_id)
+        first_token = str(
+            self.get_success(store.get_topological_token_for_event(first_event_id))
         )
 
         # Send a second message in the room, which won't be removed, and which we'll
         # use as the marker to purge events before.
         second_event_id = self.helper.send(self.room_id, "message 2")["event_id"]
-        second_token = self.get_success(
-            store.get_topological_token_for_event(second_event_id)
+        second_token = str(
+            self.get_success(store.get_topological_token_for_event(second_event_id))
         )
 
         # Send a third event in the room to ensure we don't fall under any edge case
diff --git a/tests/storage/test_purge.py b/tests/storage/test_purge.py
index 918387733b..723cd28933 100644
--- a/tests/storage/test_purge.py
+++ b/tests/storage/test_purge.py
@@ -47,8 +47,8 @@ class PurgeTests(HomeserverTestCase):
         storage = self.hs.get_storage()
 
         # Get the topological token
-        event = self.get_success(
-            store.get_topological_token_for_event(last["event_id"])
+        event = str(
+            self.get_success(store.get_topological_token_for_event(last["event_id"]))
         )
 
         # Purge everything before this topological token
@@ -74,12 +74,10 @@ class PurgeTests(HomeserverTestCase):
         storage = self.hs.get_datastore()
 
         # Set the topological token higher than it should be
-        event = self.get_success(
+        token = self.get_success(
             storage.get_topological_token_for_event(last["event_id"])
         )
-        event = "t{}-{}".format(
-            *list(map(lambda x: x + 1, map(int, event[1:].split("-"))))
-        )
+        event = "t{}-{}".format(token.topological + 1, token.stream + 1)
 
         # Purge everything before this topological token
         purge = defer.ensureDeferred(storage.purge_history(self.room_id, event, True))