Add basic tests.

synapse/storage/databases/main/event_push_actions.py

@@ -238,6 +238,13 @@ def _get_unread_counts_by_pos_txn(
     ) -> Tuple[NotifCounts, Dict[str, NotifCounts]]:
         """Get the number of unread messages for a user/room that have happened
         since the given stream ordering.
+
+        Returns:
+            A tuple of:
+                The unread messages for the main timeline
+
+                A dictionary of thread ID to unread messages for that thread.
+                Only contains threads with unread messages.
         """
 
         counts = NotifCounts()
@@ -275,7 +282,8 @@ def _get_unread_counts_by_pos_txn(
                 counts = NotifCounts(
                     notify_count=notif_count, unread_count=unread_count
                 )
-            else:
+            # TODO Delete zeroed out threads completely from the database.
+            elif notif_count or unread_count:
                 thread_counts[thread_id] = NotifCounts(
                     notify_count=notif_count, unread_count=unread_count
                 )
@@ -299,8 +307,13 @@ def _get_unread_counts_by_pos_txn(
         for highlight_count, thread_id in rows:
             if not thread_id:
                 counts.highlight_count += highlight_count
-            else:
-                thread_counts[thread_id].highlight_count += highlight_count
+            elif highlight_count:
+                if thread_id in thread_counts:
+                    thread_counts[thread_id].highlight_count += highlight_count
+                else:
+                    thread_counts[thread_id] = NotifCounts(
+                        notify_count=0, unread_count=0, highlight_count=highlight_count
+                    )
 
         # Finally we need to count push actions that aren't included in the
         # summary returned above, e.g. recent events that haven't been

tests/storage/test_event_push_actions.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 from unittest.mock import Mock
 
 from twisted.test.proto_helpers import MemoryReactor
@@ -205,6 +206,200 @@ def _mark_read(stream: int, depth: int) -> None:
         _rotate(11)
         _assert_counts(0, 0)
 
+    def test_count_aggregation_threads(self) -> None:
+        room_id = "!foo:example.com"
+        user_id = "@user1235:test"
+        thread_id = "$test7:example.com"
+
+        last_read_stream_ordering = [0]
+
+        def _assert_counts(
+            noitf_count: int,
+            highlight_count: int,
+            thread_notif_count: int,
+            thread_highlight_count: int,
+        ) -> None:
+            counts, thread_counts = self.get_success(
+                self.store.db_pool.runInteraction(
+                    "",
+                    self.store._get_unread_counts_by_pos_txn,
+                    room_id,
+                    user_id,
+                    last_read_stream_ordering[0],
+                )
+            )
+            self.assertEqual(
+                counts,
+                NotifCounts(
+                    notify_count=noitf_count,
+                    unread_count=0,  # Unread counts are tested in the sync tests.
+                    highlight_count=highlight_count,
+                ),
+            )
+            if thread_notif_count or thread_highlight_count:
+                self.assertEqual(
+                    thread_counts,
+                    {
+                        thread_id: NotifCounts(
+                            notify_count=thread_notif_count,
+                            unread_count=0,  # Unread counts are tested in the sync tests.
+                            highlight_count=thread_highlight_count,
+                        ),
+                    },
+                )
+            else:
+                self.assertEqual(thread_counts, {})
+
+        def _inject_actions(
+            stream: int, action: list, thread_id: Optional[str] = None
+        ) -> None:
+            event = Mock()
+            event.room_id = room_id
+            event.event_id = f"$test{stream}:example.com"
+            event.internal_metadata.stream_ordering = stream
+            event.internal_metadata.is_outlier.return_value = False
+            event.depth = stream
+
+            self.store._events_stream_cache.entity_has_changed(room_id, stream)
+
+            self.get_success(
+                self.store.db_pool.simple_insert(
+                    table="events",
+                    values={
+                        "stream_ordering": stream,
+                        "topological_ordering": stream,
+                        "type": "m.room.message",
+                        "room_id": room_id,
+                        "processed": True,
+                        "outlier": False,
+                        "event_id": event.event_id,
+                    },
+                )
+            )
+
+            self.get_success(
+                self.store.add_push_actions_to_staging(
+                    event.event_id,
+                    {user_id: action},
+                    False,
+                    thread_id,
+                )
+            )
+            self.get_success(
+                self.store.db_pool.runInteraction(
+                    "",
+                    self.persist_events_store._set_push_actions_for_event_and_users_txn,
+                    [(event, None)],
+                    [(event, None)],
+                )
+            )
+
+        def _rotate(stream: int) -> None:
+            self.get_success(
+                self.store.db_pool.runInteraction(
+                    "rotate-receipts", self.store._handle_new_receipts_for_notifs_txn
+                )
+            )
+
+            self.get_success(
+                self.store.db_pool.runInteraction(
+                    "rotate-notifs", self.store._rotate_notifs_before_txn, stream
+                )
+            )
+
+        def _mark_read(stream: int, depth: int) -> None:
+            last_read_stream_ordering[0] = stream
+
+            self.get_success(
+                self.store.insert_receipt(
+                    room_id,
+                    "m.read",
+                    user_id=user_id,
+                    event_ids=[f"$test{stream}:example.com"],
+                    data={},
+                )
+            )
+
+        _assert_counts(0, 0, 0, 0)
+        _inject_actions(1, PlAIN_NOTIF)
+        _assert_counts(1, 0, 0, 0)
+        _rotate(1)
+        _assert_counts(1, 0, 0, 0)
+
+        _inject_actions(2, PlAIN_NOTIF, thread_id)
+        _assert_counts(1, 0, 1, 0)
+        _rotate(2)
+        _assert_counts(1, 0, 1, 0)
+
+        _inject_actions(4, PlAIN_NOTIF)
+        _assert_counts(2, 0, 1, 0)
+        _rotate(4)
+        _assert_counts(2, 0, 1, 0)
+
+        _inject_actions(5, PlAIN_NOTIF, thread_id)
+        _assert_counts(2, 0, 2, 0)
+        _rotate(5)
+        _assert_counts(2, 0, 2, 0)
+
+        _inject_actions(6, PlAIN_NOTIF)
+        _inject_actions(7, PlAIN_NOTIF, thread_id)
+        _mark_read(5, 5)
+        _assert_counts(1, 0, 1, 0)
+
+        _mark_read(7, 7)
+        _assert_counts(0, 0, 0, 0)
+
+        _inject_actions(8, PlAIN_NOTIF)
+        _inject_actions(9, PlAIN_NOTIF, thread_id)
+        _rotate(9)
+        _assert_counts(1, 0, 1, 0)
+
+        self.get_success(
+            self.store.db_pool.simple_delete(
+                table="event_push_actions", keyvalues={"1": 1}, desc=""
+            )
+        )
+
+        _assert_counts(1, 0, 1, 0)
+
+        _mark_read(9, 9)
+        _assert_counts(0, 0, 0, 0)
+
+        _inject_actions(10, HIGHLIGHT)
+        _assert_counts(1, 1, 0, 0)
+        _rotate(10)
+        _assert_counts(1, 1, 0, 0)
+
+        _inject_actions(11, HIGHLIGHT, thread_id)
+        _assert_counts(1, 1, 1, 1)
+        _rotate(11)
+        _assert_counts(1, 1, 1, 1)
+
+        # Check that adding another notification and rotating after highlight
+        # works.
+        _inject_actions(12, PlAIN_NOTIF)
+        _rotate(12)
+        _assert_counts(2, 1, 1, 1)
+
+        _inject_actions(13, PlAIN_NOTIF, thread_id)
+        _rotate(13)
+        _assert_counts(2, 1, 2, 1)
+
+        # Check that sending read receipts at different points results in the
+        # right counts.
+        _mark_read(11, 11)
+        _assert_counts(1, 0, 1, 0)
+        _mark_read(13, 13)
+        _assert_counts(0, 0, 0, 0)
+
+        _inject_actions(14, HIGHLIGHT)
+        _inject_actions(15, HIGHLIGHT, thread_id)
+        _assert_counts(1, 1, 1, 1)
+        _mark_read(15, 15)
+        _assert_counts(0, 0, 0, 0)
+        _rotate(15)
+        _assert_counts(0, 0, 0, 0)
+
     def test_find_first_stream_ordering_after_ts(self) -> None:
         def add_event(so: int, ts: int) -> None:
             self.get_success(