Refactors migrations check. (#8814)

5 files changed, 39 insertions, 89 deletions

diff --git a/api4/channel_test.go b/api4/channel_test.go
index f871e66ea..8fd68fc08 100644
--- a/api4/channel_test.go
+++ b/api4/channel_test.go
@@ -12,8 +12,6 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-
 	"github.com/mattermost/mattermost-server/model"
 	"github.com/mattermost/mattermost-server/utils"
 )
@@ -1892,16 +1890,7 @@ func TestUpdateChannelScheme(t *testing.T) {
 
 	th.App.SetLicense(model.NewTestLicense(""))
 
-	// Mark the migration as done.
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
-
-	// Un-mark the migration at the end of the test.
-	defer func() {
-		res := <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-		assert.Nil(t, res.Err)
-	}()
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	team := &model.Team{
 		DisplayName:     "Name",
diff --git a/api4/scheme_test.go b/api4/scheme_test.go
index 461b03421..67cfda4fc 100644
--- a/api4/scheme_test.go
+++ b/api4/scheme_test.go
@@ -18,10 +18,7 @@ func TestCreateScheme(t *testing.T) {
 
 	th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
-	// Mark the migration as done.
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	// Basic test of creating a team scheme.
 	scheme1 := &model.Scheme{
@@ -137,9 +134,7 @@ func TestCreateScheme(t *testing.T) {
 	_, r6 := th.SystemAdminClient.CreateScheme(scheme6)
 	CheckNotImplementedStatus(t, r6)
 
-	// Mark the migration as not done.
-	res = <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(false)
 
 	th.LoginSystemAdmin()
 	th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
@@ -168,9 +163,7 @@ func TestGetScheme(t *testing.T) {
 		Scope:       model.SCHEME_SCOPE_TEAM,
 	}
 
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	s1, r1 := th.SystemAdminClient.CreateScheme(scheme1)
 	CheckNoError(t, r1)
@@ -210,9 +203,7 @@ func TestGetScheme(t *testing.T) {
 	_, r7 := th.Client.GetScheme(s1.Id)
 	CheckForbiddenStatus(t, r7)
 
-	// Mark the migration as not done.
-	res = <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(false)
 
 	_, r8 := th.SystemAdminClient.GetScheme(s1.Id)
 	CheckNotImplementedStatus(t, r8)
@@ -238,9 +229,7 @@ func TestGetSchemes(t *testing.T) {
 		Scope:       model.SCHEME_SCOPE_CHANNEL,
 	}
 
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	_, r1 := th.SystemAdminClient.CreateScheme(scheme1)
 	CheckNoError(t, r1)
@@ -277,9 +266,7 @@ func TestGetSchemes(t *testing.T) {
 	_, r8 := th.Client.GetSchemes("", 0, 100)
 	CheckForbiddenStatus(t, r8)
 
-	// Mark the migration as not done.
-	res = <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(false)
 
 	_, r9 := th.SystemAdminClient.GetSchemes("", 0, 100)
 	CheckNotImplementedStatus(t, r9)
@@ -291,9 +278,7 @@ func TestGetTeamsForScheme(t *testing.T) {
 
 	th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	scheme1 := &model.Scheme{
 		DisplayName: model.NewId(),
@@ -376,9 +361,7 @@ func TestGetTeamsForScheme(t *testing.T) {
 	_, ri5 := th.SystemAdminClient.GetTeamsForScheme(scheme2.Id, 0, 100)
 	CheckBadRequestStatus(t, ri5)
 
-	// Mark the migration as not done.
-	res = <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(false)
 
 	_, ri6 := th.SystemAdminClient.GetTeamsForScheme(scheme1.Id, 0, 100)
 	CheckNotImplementedStatus(t, ri6)
@@ -390,9 +373,7 @@ func TestGetChannelsForScheme(t *testing.T) {
 
 	th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	scheme1 := &model.Scheme{
 		DisplayName: model.NewId(),
@@ -477,9 +458,7 @@ func TestGetChannelsForScheme(t *testing.T) {
 	_, ri5 := th.SystemAdminClient.GetChannelsForScheme(scheme2.Id, 0, 100)
 	CheckBadRequestStatus(t, ri5)
 
-	// Mark the migration as not done.
-	res = <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(false)
 
 	_, ri6 := th.SystemAdminClient.GetChannelsForScheme(scheme1.Id, 0, 100)
 	CheckNotImplementedStatus(t, ri6)
@@ -491,10 +470,7 @@ func TestPatchScheme(t *testing.T) {
 
 	th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
-	// Mark the migration as done.
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	// Basic test of creating a team scheme.
 	scheme1 := &model.Scheme{
@@ -584,9 +560,7 @@ func TestPatchScheme(t *testing.T) {
 	_, r11 := th.SystemAdminClient.PatchScheme(s6.Id, schemePatch)
 	CheckNotImplementedStatus(t, r11)
 
-	// Mark the migration as not done.
-	res = <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	assert.Nil(t, res.Err)
+	th.App.SetPhase2PermissionsMigrationStatus(false)
 
 	th.LoginSystemAdmin()
 	th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
@@ -602,16 +576,7 @@ func TestDeleteScheme(t *testing.T) {
 	t.Run("ValidTeamScheme", func(t *testing.T) {
 		th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
-		// Mark the migration as done.
-		<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-		res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-		assert.Nil(t, res.Err)
-
-		// Un-mark the migration at the end of the test.
-		defer func() {
-			res := <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-			assert.Nil(t, res.Err)
-		}()
+		th.App.SetPhase2PermissionsMigrationStatus(true)
 
 		// Create a team scheme.
 		scheme1 := &model.Scheme{
@@ -640,7 +605,7 @@ func TestDeleteScheme(t *testing.T) {
 		assert.Zero(t, role4.DeleteAt)
 
 		// Make sure this scheme is in use by a team.
-		res = <-th.App.Srv.Store.Team().Save(&model.Team{
+		res := <-th.App.Srv.Store.Team().Save(&model.Team{
 			Name:        model.NewId(),
 			DisplayName: model.NewId(),
 			Email:       model.NewId() + "@nowhere.com",
@@ -678,16 +643,7 @@ func TestDeleteScheme(t *testing.T) {
 	t.Run("ValidChannelScheme", func(t *testing.T) {
 		th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
-		// Mark the migration as done.
-		<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-		res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-		assert.Nil(t, res.Err)
-
-		// Un-mark the migration at the end of the test.
-		defer func() {
-			res := <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-			assert.Nil(t, res.Err)
-		}()
+		th.App.SetPhase2PermissionsMigrationStatus(true)
 
 		// Create a channel scheme.
 		scheme1 := &model.Scheme{
@@ -710,7 +666,7 @@ func TestDeleteScheme(t *testing.T) {
 		assert.Zero(t, role4.DeleteAt)
 
 		// Make sure this scheme is in use by a team.
-		res = <-th.App.Srv.Store.Channel().Save(&model.Channel{
+		res := <-th.App.Srv.Store.Channel().Save(&model.Channel{
 			TeamId:      model.NewId(),
 			DisplayName: model.NewId(),
 			Name:        model.NewId(),
@@ -742,10 +698,7 @@ func TestDeleteScheme(t *testing.T) {
 	t.Run("FailureCases", func(t *testing.T) {
 		th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
-		// Mark the migration as done.
-		<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-		res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-		assert.Nil(t, res.Err)
+		th.App.SetPhase2PermissionsMigrationStatus(true)
 
 		scheme1 := &model.Scheme{
 			DisplayName: model.NewId(),
@@ -774,9 +727,7 @@ func TestDeleteScheme(t *testing.T) {
 		_, r5 := th.SystemAdminClient.DeleteScheme(s1.Id)
 		CheckNotImplementedStatus(t, r5)
 
-		// Test with migration not being done.
-		res = <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-		assert.Nil(t, res.Err)
+		th.App.SetPhase2PermissionsMigrationStatus(false)
 
 		th.App.SetLicense(model.NewTestLicense("custom_permissions_schemes"))
 
diff --git a/api4/team_test.go b/api4/team_test.go
index a3f4af0cf..f08aa6ba9 100644
--- a/api4/team_test.go
+++ b/api4/team_test.go
@@ -2059,16 +2059,7 @@ func TestUpdateTeamScheme(t *testing.T) {
 
 	th.App.SetLicense(model.NewTestLicense(""))
 
-	// Mark the migration as done.
-	<-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-	res := <-th.App.Srv.Store.System().Save(&model.System{Name: model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2, Value: "true"})
-	assert.Nil(t, res.Err)
-
-	// Un-mark the migration at the end of the test.
-	defer func() {
-		res := <-th.App.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
-		assert.Nil(t, res.Err)
-	}()
+	th.App.SetPhase2PermissionsMigrationStatus(true)
 
 	team := &model.Team{
 		DisplayName:     "Name",
diff --git a/app/app.go b/app/app.go
index d4a663e32..16470d850 100644
--- a/app/app.go
+++ b/app/app.go
@@ -92,6 +92,8 @@ type App struct {
 	clientConfig     map[string]string
 	clientConfigHash string
 	diagnosticId     string
+
+	phase2PermissionsMigrationComplete bool
 }
 
 var appCount = 0
@@ -588,3 +590,14 @@ func (a *App) DoAdvancedPermissionsMigration() {
 		mlog.Critical(fmt.Sprint(result.Err))
 	}
 }
+
+func (a *App) SetPhase2PermissionsMigrationStatus(isComplete bool) error {
+	if !isComplete {
+		res := <-a.Srv.Store.System().PermanentDeleteByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2)
+		if res.Err != nil {
+			return res.Err
+		}
+	}
+	a.phase2PermissionsMigrationComplete = isComplete
+	return nil
+}
diff --git a/app/scheme.go b/app/scheme.go
index c44690954..f070e36f8 100644
--- a/app/scheme.go
+++ b/app/scheme.go
@@ -142,10 +142,16 @@ func (a *App) GetChannelsForScheme(scheme *model.Scheme, offset int, limit int)
 }
 
 func (a *App) IsPhase2MigrationCompleted() *model.AppError {
+	if a.phase2PermissionsMigrationComplete {
+		return nil
+	}
+
 	if result := <-a.Srv.Store.System().GetByName(model.MIGRATION_KEY_ADVANCED_PERMISSIONS_PHASE_2); result.Err != nil {
 		return model.NewAppError("App.IsPhase2MigrationCompleted", "app.schemes.is_phase_2_migration_completed.not_completed.app_error", nil, result.Err.Error(), http.StatusNotImplemented)
 	}
 
+	a.phase2PermissionsMigrationComplete = true
+
 	return nil
 }