diff --git a/tests/replication/slave/storage/test_events.py b/tests/replication/slave/storage/test_events.py
index b25a06b427..eca6a443af 100644
--- a/tests/replication/slave/storage/test_events.py
+++ b/tests/replication/slave/storage/test_events.py
@@ -20,6 +20,7 @@ from synapse.api.room_versions import RoomVersions
from synapse.events import FrozenEvent, _EventInternalMetadata, make_event_from_dict
from synapse.handlers.room import RoomEventSource
from synapse.replication.slave.storage.events import SlavedEventStore
+from synapse.storage.databases.main.event_push_actions import NotifCounts
from synapse.storage.roommember import GetRoomsForUserWithStreamOrdering, RoomsForUser
from synapse.types import PersistedEventPosition
@@ -166,7 +167,7 @@ class SlavedEventStoreTestCase(BaseSlavedStoreTestCase):
self.check(
"get_unread_event_push_actions_by_room_for_user",
[ROOM_ID, USER_ID_2, event1.event_id],
- {"highlight_count": 0, "unread_count": 0, "notify_count": 0},
+ NotifCounts(highlight_count=0, unread_count=0, notify_count=0),
)
self.persist(
@@ -179,7 +180,7 @@ class SlavedEventStoreTestCase(BaseSlavedStoreTestCase):
self.check(
"get_unread_event_push_actions_by_room_for_user",
[ROOM_ID, USER_ID_2, event1.event_id],
- {"highlight_count": 0, "unread_count": 0, "notify_count": 1},
+ NotifCounts(highlight_count=0, unread_count=0, notify_count=1),
)
self.persist(
@@ -194,7 +195,7 @@ class SlavedEventStoreTestCase(BaseSlavedStoreTestCase):
self.check(
"get_unread_event_push_actions_by_room_for_user",
[ROOM_ID, USER_ID_2, event1.event_id],
- {"highlight_count": 1, "unread_count": 0, "notify_count": 2},
+ NotifCounts(highlight_count=1, unread_count=0, notify_count=2),
)
def test_get_rooms_for_user_with_stream_ordering(self):
diff --git a/tests/storage/test_event_push_actions.py b/tests/storage/test_event_push_actions.py
index bb5939ba4a..738f3ad1dc 100644
--- a/tests/storage/test_event_push_actions.py
+++ b/tests/storage/test_event_push_actions.py
@@ -14,6 +14,8 @@
from unittest.mock import Mock
+from synapse.storage.databases.main.event_push_actions import NotifCounts
+
from tests.unittest import HomeserverTestCase
USER_ID = "@user:example.com"
@@ -57,11 +59,11 @@ class EventPushActionsStoreTestCase(HomeserverTestCase):
)
self.assertEquals(
counts,
- {
- "notify_count": noitf_count,
- "unread_count": 0, # Unread counts are tested in the sync tests.
- "highlight_count": highlight_count,
- },
+ NotifCounts(
+ notify_count=noitf_count,
+ unread_count=0, # Unread counts are tested in the sync tests.
+ highlight_count=highlight_count,
+ ),
)
def _inject_actions(stream, action):
|
