Badge count fix for push notifications when CRT is enabled (#20898)

* Fix

* Fixed other cases and tests

Co-authored-by: Mattermod <mattermod@users.noreply.github.com>

store/opentracinglayer/opentracinglayer.go

@@ -11089,7 +11089,7 @@ func (s *OpenTracingLayerUserStore) GetTeamGroupUsers(teamID string) ([]*model.U
 	return result, err
 }
 
-func (s *OpenTracingLayerUserStore) GetUnreadCount(userID string) (int64, error) {
+func (s *OpenTracingLayerUserStore) GetUnreadCount(userID string, isCRTEnabled bool) (int64, error) {
 	origCtx := s.Root.Store.Context()
 	span, newCtx := tracing.StartSpanWithParentByContext(s.Root.Store.Context(), "UserStore.GetUnreadCount")
 	s.Root.Store.SetContext(newCtx)
@@ -11098,7 +11098,7 @@ func (s *OpenTracingLayerUserStore) GetUnreadCount(userID string) (int64, error)
 	}()
 
 	defer span.Finish()
-	result, err := s.UserStore.GetUnreadCount(userID)
+	result, err := s.UserStore.GetUnreadCount(userID, isCRTEnabled)
 	if err != nil {
 		span.LogFields(spanlog.Error(err))
 		ext.Error.Set(span, true)

store/retrylayer/retrylayer.go

@@ -12660,11 +12660,11 @@ func (s *RetryLayerUserStore) GetTeamGroupUsers(teamID string) ([]*model.User, e
 
 }
 
-func (s *RetryLayerUserStore) GetUnreadCount(userID string) (int64, error) {
+func (s *RetryLayerUserStore) GetUnreadCount(userID string, isCRTEnabled bool) (int64, error) {
 
 	tries := 0
 	for {
-		result, err := s.UserStore.GetUnreadCount(userID)
+		result, err := s.UserStore.GetUnreadCount(userID, isCRTEnabled)
 		if err == nil {
 			return result, nil
 		}

store/sqlstore/user_store.go

@@ -1371,9 +1371,18 @@ func (us SqlUserStore) AnalyticsActiveCountForPeriod(startTime int64, endTime in
 	return v, nil
 }
 
-func (us SqlUserStore) GetUnreadCount(userId string) (int64, error) {
+func (us SqlUserStore) GetUnreadCount(userId string, isCRTEnabled bool) (int64, error) {
+	var totalMsgCountColumn = "c.TotalMsgCount"
+	var msgCountColumn = "cm.MsgCount"
+	var mentionCountColumn = "cm.MentionCount"
+	if isCRTEnabled {
+		totalMsgCountColumn = "c.TotalMsgCountRoot"
+		msgCountColumn = "cm.MsgCountRoot"
+		mentionCountColumn = "cm.MentionCountRoot"
+	}
+
 	query := `
-		SELECT SUM(CASE WHEN c.Type = ? THEN (c.TotalMsgCount - cm.MsgCount) ELSE cm.MentionCount END)
+		SELECT SUM(CASE WHEN c.Type = ? THEN (` + totalMsgCountColumn + ` - ` + msgCountColumn + `) ELSE ` + mentionCountColumn + ` END)
 		FROM Channels c
 		INNER JOIN ChannelMembers cm
 			ON cm.ChannelId = c.Id

store/store.go

@@ -447,7 +447,7 @@ type UserStore interface {
 	PermanentDelete(userID string) error
 	AnalyticsActiveCount(timestamp int64, options model.UserCountOptions) (int64, error)
 	AnalyticsActiveCountForPeriod(startTime int64, endTime int64, options model.UserCountOptions) (int64, error)
-	GetUnreadCount(userID string) (int64, error)
+	GetUnreadCount(userID string, isCRTEnabled bool) (int64, error)
 	GetUnreadCountForChannel(userID string, channelID string) (int64, error)
 	GetAnyUnreadPostCountForChannel(userID string, channelID string) (int64, error)
 	GetRecentlyActiveUsersForTeam(teamID string, offset, limit int, viewRestrictions *model.ViewUsersRestrictions) ([]*model.User, error)

store/storetest/mocks/UserStore.go

@@ -960,20 +960,20 @@ func (_m *UserStore) GetTeamGroupUsers(teamID string) ([]*model.User, error) {
 	return r0, r1
 }
 
-// GetUnreadCount provides a mock function with given fields: userID
-func (_m *UserStore) GetUnreadCount(userID string) (int64, error) {
-	ret := _m.Called(userID)
+// GetUnreadCount provides a mock function with given fields: userID, isCRTEnabled
+func (_m *UserStore) GetUnreadCount(userID string, isCRTEnabled bool) (int64, error) {
+	ret := _m.Called(userID, isCRTEnabled)
 
 	var r0 int64
-	if rf, ok := ret.Get(0).(func(string) int64); ok {
-		r0 = rf(userID)
+	if rf, ok := ret.Get(0).(func(string, bool) int64); ok {
+		r0 = rf(userID, isCRTEnabled)
 	} else {
 		r0 = ret.Get(0).(int64)
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(userID)
+	if rf, ok := ret.Get(1).(func(string, bool) error); ok {
+		r1 = rf(userID, isCRTEnabled)
 	} else {
 		r1 = ret.Error(1)
 	}

store/storetest/user_store.go

@@ -2440,14 +2440,23 @@ func testUserUnreadCount(t *testing.T, ss store.Store) {
 	nErr = ss.Channel().IncrementMentionCount(c2.Id, []string{u2.Id}, false)
 	require.NoError(t, nErr)
 
-	badge, unreadCountErr := ss.User().GetUnreadCount(u2.Id)
+	badge, unreadCountErr := ss.User().GetUnreadCount(u2.Id, false)
 	require.NoError(t, unreadCountErr)
 	require.Equal(t, int64(3), badge, "should have 3 unread messages")
 
-	badge, unreadCountErr = ss.User().GetUnreadCount(u3.Id)
+	badge, unreadCountErr = ss.User().GetUnreadCount(u3.Id, false)
 	require.NoError(t, unreadCountErr)
 	require.Equal(t, int64(1), badge, "should have 1 unread message")
 
+	// Increment root mentions by 1
+	nErr = ss.Channel().IncrementMentionCount(c1.Id, []string{u3.Id}, true)
+	require.NoError(t, nErr)
+
+	// CRT is enabled, only root mentions are counted
+	badge, unreadCountErr = ss.User().GetUnreadCount(u3.Id, true)
+	require.NoError(t, unreadCountErr)
+	require.Equal(t, int64(1), badge, "should have 1 unread message with CRT")
+
 	badge, unreadCountErr = ss.User().GetUnreadCountForChannel(u2.Id, c1.Id)
 	require.NoError(t, unreadCountErr)
 	require.Equal(t, int64(1), badge, "should have 1 unread messages for that channel")

store/timerlayer/timerlayer.go

@@ -9983,10 +9983,10 @@ func (s *TimerLayerUserStore) GetTeamGroupUsers(teamID string) ([]*model.User, e
 	return result, err
 }
 
-func (s *TimerLayerUserStore) GetUnreadCount(userID string) (int64, error) {
+func (s *TimerLayerUserStore) GetUnreadCount(userID string, isCRTEnabled bool) (int64, error) {
 	start := time.Now()
 
-	result, err := s.UserStore.GetUnreadCount(userID)
+	result, err := s.UserStore.GetUnreadCount(userID, isCRTEnabled)
 
 	elapsed := float64(time.Since(start)) / float64(time.Second)
 	if s.Root.Metrics != nil {