upgrademanager: simplify TestMigrationFailure

This also removes `TODOTestTenantDisabled`. (I have verified the test works, albeit still flaky due to cockroachdb#106648, by temporarily removing the skip.) Release note: None
knz · Jul 17, 2023 · a2127b2 · a2127b2
1 parent e0d4e62
commit a2127b2
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 13 deletions.
diff --git a/pkg/upgrade/upgrademanager/BUILD.bazel b/pkg/upgrade/upgrademanager/BUILD.bazel
@@ -48,6 +48,7 @@ go_test(
     args = ["-test.timeout=295s"],
     deps = [
         "//pkg/base",
+        "//pkg/ccl",
         "//pkg/ccl/kvccl/kvtenantccl",
         "//pkg/clusterversion",
         "//pkg/jobs",

diff --git a/pkg/upgrade/upgrademanager/main_test.go b/pkg/upgrade/upgrademanager/main_test.go
@@ -14,6 +14,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/ccl"
 	"github.com/cockroachdb/cockroach/pkg/security/securityassets"
 	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
 	"github.com/cockroachdb/cockroach/pkg/server"
@@ -25,6 +26,7 @@ func TestMain(m *testing.M) {
 	securityassets.SetLoader(securitytest.EmbeddedAssets)
 	serverutils.InitTestServerFactory(server.TestServerFactory)
 	serverutils.InitTestClusterFactory(testcluster.TestClusterFactory)
+	defer ccl.TestingEnableEnterprise()()
 	os.Exit(m.Run())
 }
 

diff --git a/pkg/upgrade/upgrademanager/manager_external_test.go b/pkg/upgrade/upgrademanager/manager_external_test.go
@@ -689,24 +689,20 @@ func TestMigrationFailure(t *testing.T) {
 	fenceVersion := upgrade.FenceVersionFor(ctx, clusterversion.ClusterVersion{Version: failVersion}).Version
 	t.Logf("test will fail at version: %s", failVersion.String())
 
-	// Create a storage cluster for the tenant
-	testCluster := serverutils.StartNewTestCluster(t, 1, base.TestClusterArgs{
-		ServerArgs: base.TestServerArgs{
-			DefaultTestTenant: base.TODOTestTenantDisabled,
-			Knobs: base.TestingKnobs{
-				SQLEvalContext: &eval.TestingKnobs{
-					TenantLogicalVersionKeyOverride: startVersionKey,
-				},
+	// Create a storage cluster for the tenant.
+	s, goDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+		DefaultTestTenant: base.TestControlsTenantsExplicitly,
+		Knobs: base.TestingKnobs{
+			SQLEvalContext: &eval.TestingKnobs{
+				TenantLogicalVersionKeyOverride: startVersionKey,
 			},
 		},
 	})
-	defer testCluster.Stopper().Stop(ctx)
+	defer s.Stopper().Stop(ctx)
 
 	// Set the version override so that the tenant is able to upgrade. If this is
 	// not set, the tenant treats the storage cluster as if it had the oldest
 	// supported binary version.
-	s := testCluster.Server(0)
-	goDB := serverutils.OpenDBConn(t, s.ServingSQLAddr(), "system", false, s.Stopper())
 	_, err := goDB.Exec(`ALTER TENANT ALL SET CLUSTER SETTING version = $1`, endVersion.String())
 	require.NoError(t, err)
 
@@ -722,7 +718,7 @@ func TestMigrationFailure(t *testing.T) {
 		false,
 	)
 	require.NoError(t, clusterversion.Initialize(ctx, startVersion, &tenantSettings.SV))
-	tenant, db := serverutils.StartTenant(t, testCluster.Server(0), base.TestTenantArgs{
+	tenant, db := serverutils.StartTenant(t, s, base.TestTenantArgs{
 		TenantID: roachpb.MustMakeTenantID(10),
 		Settings: tenantSettings,
 		TestingKnobs: base.TestingKnobs{
@@ -773,7 +769,7 @@ func TestMigrationFailure(t *testing.T) {
 	checkActiveVersion(t, startVersion)
 	checkSettingVersion(t, startVersion)
 
-	// Try to finalize
+	// Try to finalize.
 	_, err = db.Exec(`SET CLUSTER SETTING version = $1`, endVersion.String())
 	require.Error(t, err)
 	checkActiveVersion(t, fenceVersion)