feat(unified): data migration integration tests (#114418)

* feat: unified storage migrations integration tests

* chore: add comment and adjust db path name

* chore: refactor test cases into interface

pkg/tests/apis/helper.go

@@ -93,7 +93,18 @@ type K8sTestHelper struct {
 	userSvc user.Service
 }
 
+type K8sTestHelperOpts struct {
+	testinfra.GrafanaOpts
+	// If provided, these users will be used instead of creating new ones
+	Org1Users *OrgUsers
+	OrgBUsers *OrgUsers
+}
+
 func NewK8sTestHelper(t *testing.T, opts testinfra.GrafanaOpts) *K8sTestHelper {
+	return NewK8sTestHelperWithOpts(t, K8sTestHelperOpts{GrafanaOpts: opts})
+}
+
+func NewK8sTestHelperWithOpts(t *testing.T, opts K8sTestHelperOpts) *K8sTestHelper {
 	t.Helper()
 
 	// Use GRPC server when not configured
@@ -111,9 +122,12 @@ func NewK8sTestHelper(t *testing.T, opts testinfra.GrafanaOpts) *K8sTestHelper {
 		path = opts.DirPath
 	)
 	if opts.Dir == "" && opts.DirPath == "" {
-		dir, path = testinfra.CreateGrafDir(t, opts)
+		dir, path = testinfra.CreateGrafDir(t, opts.GrafanaOpts)
+	}
+	listenerAddress, env, testDB := testinfra.StartGrafanaEnvWithDB(t, dir, path)
+	if !opts.DisableDBCleanup {
+		t.Cleanup(testDB.Cleanup)
 	}
-	listenerAddress, env := testinfra.StartGrafanaEnv(t, dir, path)
 
 	c := &K8sTestHelper{
 		env:             *env,
@@ -143,8 +157,24 @@ func NewK8sTestHelper(t *testing.T, opts testinfra.GrafanaOpts) *K8sTestHelper {
 	_ = c.CreateOrg(Org1)
 	_ = c.CreateOrg(Org2)
 
-	c.Org1 = c.createTestUsers(Org1)
-	c.OrgB = c.createTestUsers(Org2)
+	if opts.Org1Users != nil {
+		c.Org1 = *opts.Org1Users
+		c.Org1.Admin.baseURL = listenerAddress
+		c.Org1.Editor.baseURL = listenerAddress
+		c.Org1.Viewer.baseURL = listenerAddress
+		c.Org1.None.baseURL = listenerAddress
+	} else {
+		c.Org1 = c.createTestUsers(Org1)
+	}
+	if opts.OrgBUsers != nil {
+		c.OrgB = *opts.OrgBUsers
+		c.OrgB.Admin.baseURL = listenerAddress
+		c.OrgB.Editor.baseURL = listenerAddress
+		c.OrgB.Viewer.baseURL = listenerAddress
+		c.OrgB.None.baseURL = listenerAddress
+	} else {
+		c.OrgB = c.createTestUsers(Org2)
+	}
 
 	c.loadAPIGroups()
 

pkg/tests/testinfra/testinfra.go

@@ -49,6 +49,12 @@ func StartGrafana(t *testing.T, grafDir, cfgPath string) (string, db.DB) {
 }
 
 func StartGrafanaEnv(t *testing.T, grafDir, cfgPath string) (string, *server.TestEnv) {
+	addr, env, testDB := StartGrafanaEnvWithDB(t, grafDir, cfgPath)
+	t.Cleanup(testDB.Cleanup)
+	return addr, env
+}
+
+func StartGrafanaEnvWithDB(t *testing.T, grafDir, cfgPath string) (string, *server.TestEnv, *sqlutil.TestDB) {
 	t.Helper()
 	ctx := context.Background()
 
@@ -93,7 +99,6 @@ func StartGrafanaEnv(t *testing.T, grafDir, cfgPath string) (string, *server.Tes
 	// Use proper database type based on the environment variable GRAFANA_TEST_DB in tests
 	testDB, err := sqlutil.GetTestDB(sqlutil.GetTestDBType())
 	require.NoError(t, err)
-	t.Cleanup(testDB.Cleanup)
 
 	dbCfg := cfg.Raw.Section("database")
 	dbCfg.Key("type").SetValue(testDB.DriverName)
@@ -169,7 +174,7 @@ func StartGrafanaEnv(t *testing.T, grafDir, cfgPath string) (string, *server.Tes
 
 	t.Logf("Grafana is listening on %s", addr)
 
-	return addr, env
+	return addr, env, testDB
 }
 
 // CreateGrafDir creates the Grafana directory.
@@ -538,6 +543,12 @@ func CreateGrafDir(t *testing.T, opts GrafanaOpts) (string, string) {
 		_, err = section.NewKey("max_page_size_bytes", fmt.Sprintf("%d", opts.UnifiedStorageMaxPageSizeBytes))
 		require.NoError(t, err)
 	}
+	if opts.DisableDataMigrations {
+		section, err := getOrCreateSection("unified_storage")
+		require.NoError(t, err)
+		_, err = section.NewKey("disable_data_migrations", "true")
+		require.NoError(t, err)
+	}
 	if opts.PermittedProvisioningPaths != "" {
 		_, err = pathsSect.NewKey("permitted_provisioning_paths", opts.PermittedProvisioningPaths)
 		require.NoError(t, err)
@@ -637,6 +648,8 @@ type GrafanaOpts struct {
 	EnableSCIM                            bool
 	APIServerRuntimeConfig                string
 	DisableControllers                    bool
+	DisableDBCleanup                      bool
+	DisableDataMigrations                 bool
 	SecretsManagerEnableDBMigrations      bool
 
 	// Allow creating grafana dir beforehand