[MM-15296] Migrate "Preference.Save" to Sync by default (#10866)

* response format changes

* Generated mocks
Fixed all references of Preferences.Save

* Remove old code from store.go (incorrect merge)

* Review change - Add validations on count and error

* Review change - return from root level

* Review change - return 0 as nil value for int

* fix initialisation of err in preference_store

app/channel.go

@@ -1059,8 +1059,8 @@ func (a *App) AddDirectChannels(teamId string, user *model.User) *model.AppError
 		}
 	}
 
-	if result := <-a.Srv.Store.Preference().Save(&preferences); result.Err != nil {
-		return model.NewAppError("AddDirectChannels", "api.user.add_direct_channels_and_forget.failed.error", map[string]interface{}{"UserId": user.Id, "TeamId": teamId, "Error": result.Err.Error()}, "", http.StatusInternalServerError)
+	if _, err := a.Srv.Store.Preference().Save(&preferences); err != nil {
+		return model.NewAppError("AddDirectChannels", "api.user.add_direct_channels_and_forget.failed.error", map[string]interface{}{"UserId": user.Id, "TeamId": teamId, "Error": err.Error()}, "", http.StatusInternalServerError)
 	}
 
 	return nil

app/command_expand_collapse.go

@@ -68,7 +68,7 @@ func (a *App) setCollapsePreference(args *model.CommandArgs, isCollapse bool) *m
 		Value:    strconv.FormatBool(isCollapse),
 	}
 
-	if result := <-a.Srv.Store.Preference().Save(&model.Preferences{pref}); result.Err != nil {
+	if _, err := a.Srv.Store.Preference().Save(&model.Preferences{pref}); err != nil {
 		return &model.CommandResponse{Text: args.T("api.command_expand_collapse.fail.app_error"), ResponseType: model.COMMAND_RESPONSE_TYPE_EPHEMERAL}
 	}
 

app/email_batching_test.go

@@ -115,12 +115,14 @@ func TestCheckPendingNotifications(t *testing.T) {
 	channelMember.LastViewedAt = 9999999
 	store.Must(th.App.Srv.Store.Channel().UpdateMember(channelMember))
 
-	store.Must(th.App.Srv.Store.Preference().Save(&model.Preferences{{
+	count, err := th.App.Srv.Store.Preference().Save(&model.Preferences{{
 		UserId:   th.BasicUser.Id,
 		Category: model.PREFERENCE_CATEGORY_NOTIFICATIONS,
 		Name:     model.PREFERENCE_NAME_EMAIL_INTERVAL,
 		Value:    "60",
-	}}))
+	}})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	// test that notifications aren't sent before interval
 	job.checkPendingNotifications(time.Unix(10001, 0), func(string, []*batchedNotification) {})
@@ -256,12 +258,14 @@ func TestCheckPendingNotificationsCantParseInterval(t *testing.T) {
 	store.Must(th.App.Srv.Store.Channel().UpdateMember(channelMember))
 
 	// preference value is not an integer, so we'll fall back to the default 15min value
-	store.Must(th.App.Srv.Store.Preference().Save(&model.Preferences{{
+	count, err := th.App.Srv.Store.Preference().Save(&model.Preferences{{
 		UserId:   th.BasicUser.Id,
 		Category: model.PREFERENCE_CATEGORY_NOTIFICATIONS,
 		Name:     model.PREFERENCE_NAME_EMAIL_INTERVAL,
 		Value:    "notAnIntegerValue",
-	}}))
+	}})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	job.pendingNotifications[th.BasicUser.Id] = []*batchedNotification{
 		{

app/export_test.go

@@ -10,7 +10,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/mattermost/mattermost-server/model"
-	"github.com/mattermost/mattermost-server/store"
 )
 
 func TestReactionsOfPost(t *testing.T) {
@@ -81,7 +80,10 @@ func TestExportUserChannels(t *testing.T) {
 	}
 	var preferences model.Preferences
 	preferences = append(preferences, preference)
-	store.Must(th.App.Srv.Store.Preference().Save(&preferences))
+	count, err := th.App.Srv.Store.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
+
 	th.App.UpdateChannelMemberNotifyProps(notifyProps, channel.Id, user.Id)
 	exportData, err := th.App.buildUserChannelMemberships(user.Id, team.Id)
 	require.Nil(t, err)

app/import_functions.go

@@ -636,8 +636,8 @@ func (a *App) ImportUser(data *UserImportData, dryRun bool) *model.AppError {
 	}
 
 	if len(preferences) > 0 {
-		if result := <-a.Srv.Store.Preference().Save(&preferences); result.Err != nil {
-			return model.NewAppError("BulkImport", "app.import.import_user.save_preferences.error", nil, result.Err.Error(), http.StatusInternalServerError)
+		if _, err := a.Srv.Store.Preference().Save(&preferences); err != nil {
+			return model.NewAppError("BulkImport", "app.import.import_user.save_preferences.error", nil, err.Error(), http.StatusInternalServerError)
 		}
 	}
 
@@ -721,8 +721,8 @@ func (a *App) ImportUserTeams(user *model.User, data *[]UserTeamImportData) *mod
 	}
 
 	if len(teamThemePreferences) > 0 {
-		if result := <-a.Srv.Store.Preference().Save(&teamThemePreferences); result.Err != nil {
-			return model.NewAppError("BulkImport", "app.import.import_user_teams.save_preferences.error", nil, result.Err.Error(), http.StatusInternalServerError)
+		if _, err := a.Srv.Store.Preference().Save(&teamThemePreferences); err != nil {
+			return model.NewAppError("BulkImport", "app.import.import_user_teams.save_preferences.error", nil, err.Error(), http.StatusInternalServerError)
 		}
 	}
 
@@ -818,8 +818,8 @@ func (a *App) ImportUserChannels(user *model.User, team *model.Team, teamMember
 	}
 
 	if len(preferences) > 0 {
-		if result := <-a.Srv.Store.Preference().Save(&preferences); result.Err != nil {
-			return model.NewAppError("BulkImport", "app.import.import_user_channels.save_preferences.error", nil, result.Err.Error(), http.StatusInternalServerError)
+		if _, err := a.Srv.Store.Preference().Save(&preferences); err != nil {
+			return model.NewAppError("BulkImport", "app.import.import_user_channels.save_preferences.error", nil, err.Error(), http.StatusInternalServerError)
 		}
 	}
 
@@ -1024,8 +1024,8 @@ func (a *App) ImportPost(data *PostImportData, dryRun bool) *model.AppError {
 		}
 
 		if len(preferences) > 0 {
-			if result := <-a.Srv.Store.Preference().Save(&preferences); result.Err != nil {
-				return model.NewAppError("BulkImport", "app.import.import_post.save_preferences.error", nil, result.Err.Error(), http.StatusInternalServerError)
+			if _, err := a.Srv.Store.Preference().Save(&preferences); err != nil {
+				return model.NewAppError("BulkImport", "app.import.import_post.save_preferences.error", nil, err.Error(), http.StatusInternalServerError)
 			}
 		}
 	}
@@ -1129,9 +1129,9 @@ func (a *App) ImportDirectChannel(data *DirectChannelImportData, dryRun bool) *m
 		}
 	}
 
-	if result := <-a.Srv.Store.Preference().Save(&preferences); result.Err != nil {
-		result.Err.StatusCode = http.StatusBadRequest
-		return result.Err
+	if _, err := a.Srv.Store.Preference().Save(&preferences); err != nil {
+		err.StatusCode = http.StatusBadRequest
+		return err
 	}
 
 	if data.Header != nil {
@@ -1248,8 +1248,8 @@ func (a *App) ImportDirectPost(data *DirectPostImportData, dryRun bool) *model.A
 		}
 
 		if len(preferences) > 0 {
-			if result := <-a.Srv.Store.Preference().Save(&preferences); result.Err != nil {
-				return model.NewAppError("BulkImport", "app.import.import_direct_post.save_preferences.error", nil, result.Err.Error(), http.StatusInternalServerError)
+			if _, err := a.Srv.Store.Preference().Save(&preferences); err != nil {
+				return model.NewAppError("BulkImport", "app.import.import_direct_post.save_preferences.error", nil, err.Error(), http.StatusInternalServerError)
 			}
 		}
 	}

app/oauth.go

@@ -186,8 +186,8 @@ func (a *App) AllowOAuthAppAccessToUser(userId string, authRequest *model.Author
 		Value:    authRequest.Scope,
 	}
 
-	if result = <-a.Srv.Store.Preference().Save(&model.Preferences{authorizedApp}); result.Err != nil {
-		mlog.Error(result.Err.Error())
+	if _, err = a.Srv.Store.Preference().Save(&model.Preferences{authorizedApp}); err != nil {
+		mlog.Error(err.Error())
 		return authRequest.RedirectUri + "?error=server_error&state=" + authRequest.State, nil
 	}
 

app/preference.go

@@ -48,9 +48,9 @@ func (a *App) UpdatePreferences(userId string, preferences model.Preferences) *m
 		}
 	}
 
-	if result := <-a.Srv.Store.Preference().Save(&preferences); result.Err != nil {
-		result.Err.StatusCode = http.StatusBadRequest
-		return result.Err
+	if _, err := a.Srv.Store.Preference().Save(&preferences); err != nil {
+		err.StatusCode = http.StatusBadRequest
+		return err
 	}
 
 	message := model.NewWebSocketEvent(model.WEBSOCKET_EVENT_PREFERENCES_CHANGED, "", "", userId, nil)

app/user.go

@@ -310,8 +310,8 @@ func (a *App) createUser(user *model.User) (*model.User, *model.AppError) {
 	}
 
 	pref := model.Preference{UserId: ruser.Id, Category: model.PREFERENCE_CATEGORY_TUTORIAL_STEPS, Name: ruser.Id, Value: "0"}
-	if presult := <-a.Srv.Store.Preference().Save(&model.Preferences{pref}); presult.Err != nil {
-		mlog.Error(fmt.Sprintf("Encountered error saving tutorial preference, err=%v", presult.Err.Message))
+	if _, err := a.Srv.Store.Preference().Save(&model.Preferences{pref}); err != nil {
+		mlog.Error(fmt.Sprintf("Encountered error saving tutorial preference, err=%v", err.Message))
 	}
 
 	ruser.Sanitize(map[string]bool{})

store/sqlstore/preference_store.go

@@ -54,31 +54,25 @@ func (s SqlPreferenceStore) DeleteUnusedFeatures() {
 	s.GetMaster().Exec(sql, queryParams)
 }
 
-func (s SqlPreferenceStore) Save(preferences *model.Preferences) store.StoreChannel {
-	return store.Do(func(result *store.StoreResult) {
-		// wrap in a transaction so that if one fails, everything fails
-		transaction, err := s.GetMaster().Begin()
-		if err != nil {
-			result.Err = model.NewAppError("SqlPreferenceStore.Save", "store.sql_preference.save.open_transaction.app_error", nil, err.Error(), http.StatusInternalServerError)
-		} else {
-			defer finalizeTransaction(transaction)
-			for _, preference := range *preferences {
-				if upsertResult := s.save(transaction, &preference); upsertResult.Err != nil {
-					*result = upsertResult
-					break
-				}
-			}
+func (s SqlPreferenceStore) Save(preferences *model.Preferences) (int, *model.AppError) {
+	// wrap in a transaction so that if one fails, everything fails
+	transaction, err := s.GetMaster().Begin()
+	if err != nil {
+		return 0, model.NewAppError("SqlPreferenceStore.Save", "store.sql_preference.save.open_transaction.app_error", nil, err.Error(), http.StatusInternalServerError)
+	}
 
-			if result.Err == nil {
-				if err := transaction.Commit(); err != nil {
-					// don't need to rollback here since the transaction is already closed
-					result.Err = model.NewAppError("SqlPreferenceStore.Save", "store.sql_preference.save.commit_transaction.app_error", nil, err.Error(), http.StatusInternalServerError)
-				} else {
-					result.Data = len(*preferences)
-				}
-			}
+	defer finalizeTransaction(transaction)
+	for _, preference := range *preferences {
+		if upsertResult := s.save(transaction, &preference); upsertResult.Err != nil {
+			return 0, upsertResult.Err
 		}
-	})
+	}
+
+	if err := transaction.Commit(); err != nil {
+		// don't need to rollback here since the transaction is already closed
+		return 0, model.NewAppError("SqlPreferenceStore.Save", "store.sql_preference.save.commit_transaction.app_error", nil, err.Error(), http.StatusInternalServerError)
+	}
+	return len(*preferences), nil
 }
 
 func (s SqlPreferenceStore) save(transaction *gorp.Transaction, preference *model.Preference) store.StoreResult {

store/sqlstore/preference_store_test.go

@@ -9,6 +9,8 @@ import (
 	"github.com/mattermost/mattermost-server/model"
 	"github.com/mattermost/mattermost-server/store"
 	"github.com/mattermost/mattermost-server/store/storetest"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestPreferenceStore(t *testing.T) {
@@ -50,7 +52,9 @@ func TestDeleteUnusedFeatures(t *testing.T) {
 			},
 		}
 
-		store.Must(ss.Preference().Save(&features))
+		count, err := ss.Preference().Save(&features)
+		require.Nil(t, err)
+		require.Equal(t, 4, count)
 
 		ss.Preference().(*SqlPreferenceStore).DeleteUnusedFeatures()
 

store/store.go

@@ -428,7 +428,7 @@ type CommandWebhookStore interface {
 }
 
 type PreferenceStore interface {
-	Save(preferences *model.Preferences) StoreChannel
+	Save(preferences *model.Preferences) (int, *model.AppError)
 	GetCategory(userId string, category string) (model.Preferences, *model.AppError)
 	Get(userId string, category string, name string) (*model.Preference, *model.AppError)
 	GetAll(userId string) StoreChannel

store/storetest/mocks/PreferenceStore.go

@@ -183,17 +183,24 @@ func (_m *PreferenceStore) PermanentDeleteByUser(userId string) *model.AppError
 }
 
 // Save provides a mock function with given fields: preferences
-func (_m *PreferenceStore) Save(preferences *model.Preferences) store.StoreChannel {
+func (_m *PreferenceStore) Save(preferences *model.Preferences) (int, *model.AppError) {
 	ret := _m.Called(preferences)
 
-	var r0 store.StoreChannel
-	if rf, ok := ret.Get(0).(func(*model.Preferences) store.StoreChannel); ok {
+	var r0 int
+	if rf, ok := ret.Get(0).(func(*model.Preferences) int); ok {
 		r0 = rf(preferences)
 	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(store.StoreChannel)
+		r0 = ret.Get(0).(int)
+	}
+
+	var r1 *model.AppError
+	if rf, ok := ret.Get(1).(func(*model.Preferences) *model.AppError); ok {
+		r1 = rf(preferences)
+	} else {
+		if ret.Get(1) != nil {
+			r1 = ret.Get(1).(*model.AppError)
 		}
 	}
 
-	return r0
+	return r0, r1
 }

store/storetest/oauth_store.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/mattermost/mattermost-server/model"
 	"github.com/mattermost/mattermost-server/store"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestOAuthStore(t *testing.T, ss store.Store) {
@@ -333,7 +335,9 @@ func testOAuthGetAuthorizedApps(t *testing.T, ss store.Store) {
 	p.Category = model.PREFERENCE_CATEGORY_AUTHORIZED_OAUTH_APP
 	p.Name = a1.Id
 	p.Value = "true"
-	store.Must(ss.Preference().Save(&model.Preferences{p}))
+	count, err := ss.Preference().Save(&model.Preferences{p})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	if result := <-ss.OAuth().GetAuthorizedApps(a1.CreatorId, 0, 1000); result.Err != nil {
 		t.Fatal(result.Err)
@@ -359,7 +363,9 @@ func testOAuthGetAccessDataByUserForApp(t *testing.T, ss store.Store) {
 	p.Category = model.PREFERENCE_CATEGORY_AUTHORIZED_OAUTH_APP
 	p.Name = a1.Id
 	p.Value = "true"
-	store.Must(ss.Preference().Save(&model.Preferences{p}))
+	count, err := ss.Preference().Save(&model.Preferences{p})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	if result := <-ss.OAuth().GetAuthorizedApps(a1.CreatorId, 0, 1000); result.Err != nil {
 		t.Fatal(result.Err)

store/storetest/post_store.go

@@ -1389,7 +1389,9 @@ func testPostStoreGetFlaggedPostsForTeam(t *testing.T, ss store.Store, s SqlSupp
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err := ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r2 := (<-ss.Post().GetFlaggedPostsForTeam(o1.UserId, c1.TeamId, 0, 2)).Data.(*model.PostList)
 
@@ -1406,7 +1408,9 @@ func testPostStoreGetFlaggedPostsForTeam(t *testing.T, ss store.Store, s SqlSupp
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err = ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r3 := (<-ss.Post().GetFlaggedPostsForTeam(o1.UserId, c1.TeamId, 0, 1)).Data.(*model.PostList)
 
@@ -1441,7 +1445,9 @@ func testPostStoreGetFlaggedPostsForTeam(t *testing.T, ss store.Store, s SqlSupp
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err = ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r4 = (<-ss.Post().GetFlaggedPostsForTeam(o1.UserId, c1.TeamId, 0, 2)).Data.(*model.PostList)
 
@@ -1457,7 +1463,9 @@ func testPostStoreGetFlaggedPostsForTeam(t *testing.T, ss store.Store, s SqlSupp
 			Value:    "true",
 		},
 	}
-	store.Must(ss.Preference().Save(&preferences))
+	count, err = ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r4 = (<-ss.Post().GetFlaggedPostsForTeam(o1.UserId, c1.TeamId, 0, 2)).Data.(*model.PostList)
 
@@ -1479,7 +1487,9 @@ func testPostStoreGetFlaggedPostsForTeam(t *testing.T, ss store.Store, s SqlSupp
 			Value:    "true",
 		},
 	}
-	store.Must(ss.Preference().Save(&preferences))
+	count, err = ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r4 = (<-ss.Post().GetFlaggedPostsForTeam(o1.UserId, c1.TeamId, 0, 10)).Data.(*model.PostList)
 
@@ -1529,7 +1539,9 @@ func testPostStoreGetFlaggedPosts(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err := ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r2 := (<-ss.Post().GetFlaggedPosts(o1.UserId, 0, 2)).Data.(*model.PostList)
 
@@ -1546,7 +1558,9 @@ func testPostStoreGetFlaggedPosts(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err = ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r3 := (<-ss.Post().GetFlaggedPosts(o1.UserId, 0, 1)).Data.(*model.PostList)
 
@@ -1581,7 +1595,9 @@ func testPostStoreGetFlaggedPosts(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err = ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r4 = (<-ss.Post().GetFlaggedPosts(o1.UserId, 0, 2)).Data.(*model.PostList)
 
@@ -1634,7 +1650,9 @@ func testPostStoreGetFlaggedPostsForChannel(t *testing.T, ss store.Store) {
 		Value:    "true",
 	}
 
-	store.Must(ss.Preference().Save(&model.Preferences{preference}))
+	count, err := ss.Preference().Save(&model.Preferences{preference})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r = (<-ss.Post().GetFlaggedPostsForChannel(o1.UserId, o1.ChannelId, 0, 10)).Data.(*model.PostList)
 
@@ -1643,10 +1661,14 @@ func testPostStoreGetFlaggedPostsForChannel(t *testing.T, ss store.Store) {
 	}
 
 	preference.Name = o2.Id
-	store.Must(ss.Preference().Save(&model.Preferences{preference}))
+	count, err = ss.Preference().Save(&model.Preferences{preference})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	preference.Name = o3.Id
-	store.Must(ss.Preference().Save(&model.Preferences{preference}))
+	count, err = ss.Preference().Save(&model.Preferences{preference})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r = (<-ss.Post().GetFlaggedPostsForChannel(o1.UserId, o1.ChannelId, 0, 1)).Data.(*model.PostList)
 
@@ -1673,7 +1695,9 @@ func testPostStoreGetFlaggedPostsForChannel(t *testing.T, ss store.Store) {
 	}
 
 	preference.Name = o4.Id
-	store.Must(ss.Preference().Save(&model.Preferences{preference}))
+	count, err = ss.Preference().Save(&model.Preferences{preference})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	r = (<-ss.Post().GetFlaggedPostsForChannel(o1.UserId, o4.ChannelId, 0, 10)).Data.(*model.PostList)
 

store/storetest/preference_store.go

@@ -4,6 +4,7 @@
 package storetest
 
 import (
+	"github.com/stretchr/testify/require"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -42,8 +43,10 @@ func testPreferenceSave(t *testing.T, ss store.Store) {
 			Value:    "value1b",
 		},
 	}
-	if count := store.Must(ss.Preference().Save(&preferences)); count != 2 {
+	if count, err := ss.Preference().Save(&preferences); count != 2 {
 		t.Fatal("got incorrect number of rows saved")
+	} else if err != nil {
+		t.Fatal("saving preference returned error")
 	}
 
 	for _, preference := range preferences {
@@ -54,8 +57,10 @@ func testPreferenceSave(t *testing.T, ss store.Store) {
 
 	preferences[0].Value = "value2a"
 	preferences[1].Value = "value2b"
-	if count := store.Must(ss.Preference().Save(&preferences)); count != 2 {
+	if count, err := ss.Preference().Save(&preferences); count != 2 {
 		t.Fatal("got incorrect number of rows saved")
+	} else if err != nil {
+		t.Fatal("saving preference returned error")
 	}
 
 	for _, preference := range preferences {
@@ -93,7 +98,9 @@ func testPreferenceGet(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err := ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 4, count)
 
 	if data, err := ss.Preference().Get(userId, category, name); err != nil {
 		t.Fatal(err)
@@ -138,7 +145,9 @@ func testPreferenceGetCategory(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err := ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 4, count)
 
 	if preferencesByCategory, err := ss.Preference().GetCategory(userId, category); err != nil {
 		t.Fatal(err)
@@ -187,7 +196,9 @@ func testPreferenceGetAll(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err := ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 4, count)
 
 	if result := <-ss.Preference().GetAll(userId); result.Err != nil {
 		t.Fatal(result.Err)
@@ -233,7 +244,9 @@ func testPreferenceDeleteByUser(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&preferences))
+	count, err := ss.Preference().Save(&preferences)
+	require.Nil(t, err)
+	require.Equal(t, 4, count)
 
 	if err := ss.Preference().PermanentDeleteByUser(userId); err != nil {
 		t.Fatal(err)
@@ -281,7 +294,9 @@ func testIsFeatureEnabled(t *testing.T, ss store.Store) {
 		},
 	}
 
-	store.Must(ss.Preference().Save(&features))
+	count, err := ss.Preference().Save(&features)
+	require.Nil(t, err)
+	require.Equal(t, 5, count)
 
 	if result := <-ss.Preference().IsFeatureEnabled(feature1, userId); result.Err != nil {
 		t.Fatal(result.Err)
@@ -318,7 +333,9 @@ func testPreferenceDelete(t *testing.T, ss store.Store) {
 		Value:    "value1a",
 	}
 
-	store.Must(ss.Preference().Save(&model.Preferences{preference}))
+	count, err := ss.Preference().Save(&model.Preferences{preference})
+	require.Nil(t, err)
+	require.Equal(t, 1, count)
 
 	if prefs := store.Must(ss.Preference().GetAll(preference.UserId)).(model.Preferences); len([]model.Preference(prefs)) != 1 {
 		t.Fatal("should've returned 1 preference")
@@ -351,7 +368,9 @@ func testPreferenceDeleteCategory(t *testing.T, ss store.Store) {
 		Value:    "value1a",
 	}
 
-	store.Must(ss.Preference().Save(&model.Preferences{preference1, preference2}))
+	count, err := ss.Preference().Save(&model.Preferences{preference1, preference2})
+	require.Nil(t, err)
+	require.Equal(t, 2, count)
 
 	if prefs := store.Must(ss.Preference().GetAll(userId)).(model.Preferences); len([]model.Preference(prefs)) != 2 {
 		t.Fatal("should've returned 2 preferences")
@@ -386,7 +405,9 @@ func testPreferenceDeleteCategoryAndName(t *testing.T, ss store.Store) {
 		Value:    "value1a",
 	}
 
-	store.Must(ss.Preference().Save(&model.Preferences{preference1, preference2}))
+	count, err := ss.Preference().Save(&model.Preferences{preference1, preference2})
+	require.Nil(t, err)
+	require.Equal(t, 2, count)
 
 	if prefs := store.Must(ss.Preference().GetAll(userId)).(model.Preferences); len([]model.Preference(prefs)) != 1 {
 		t.Fatal("should've returned 1 preference")
@@ -434,9 +455,11 @@ func testPreferenceCleanupFlagsBatch(t *testing.T, ss store.Store) {
 		Value:    "true",
 	}
 
-	store.Must(ss.Preference().Save(&model.Preferences{preference1, preference2}))
+	count, err := ss.Preference().Save(&model.Preferences{preference1, preference2})
+	require.Nil(t, err)
+	require.Equal(t, 2, count)
 
-	_, err := ss.Preference().CleanupFlagsBatch(10000)
+	_, err = ss.Preference().CleanupFlagsBatch(10000)
 	assert.Nil(t, err)
 
 	_, err = ss.Preference().Get(userId, category, preference1.Name)