kvserver: fix rebalance obj test on arm mac

kvoli · kvoli · commit 4c2d4c39b711 · 2023-02-10T20:29:57.000Z
Previosly,`TestRebalanceObjectiveManager` and `TestLoadBasedRebalancingObjective` would assert assuming that it was possible for the test host to use `grunning`, however this is not true for ARM Mac. This patch ammends these tests to test a subset of behavior when `grunning` isn't supported and then exit. Fixes: #96934 Release note: None
diff --git a/pkg/kv/kvserver/BUILD.bazel b/pkg/kv/kvserver/BUILD.bazel
@@ -430,6 +430,7 @@ go_test(
         "//pkg/util/ctxgroup",
         "//pkg/util/encoding",
         "//pkg/util/errorutil",
+        "//pkg/util/grunning",
         "//pkg/util/hlc",
         "//pkg/util/humanizeutil",
         "//pkg/util/leaktest",
diff --git a/pkg/kv/kvserver/rebalance_objective_test.go b/pkg/kv/kvserver/rebalance_objective_test.go
@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/util/grunning"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
@@ -82,6 +83,29 @@ func TestLoadBasedRebalancingObjective(t *testing.T) {
 	defer log.Scope(t).Close(t)
 	ctx := context.Background()
 
+	// On ARM MacOS and other architectures, grunning isn't supported so
+	// changing the objective to CPU should never work. If this test is  run on
+	// one of these unsupported aarch, test this behavior only.
+	if !grunning.Supported() {
+		st := cluster.MakeTestingClusterSettings()
+
+		gossipStoreDescProvider := testMakeProviderNotifier(allPositiveCPUMap)
+		LoadBasedRebalancingObjective.Override(ctx, &st.SV, int64(LBRebalancingQueries))
+		require.Equal(t,
+			LBRebalancingQueries,
+			ResolveLBRebalancingObjective(ctx, st, gossipStoreDescProvider.GetStores()),
+		)
+
+		// Despite setting to CPU, only QPS should be returned since this aarch
+		// doesn't support grunning.
+		LoadBasedRebalancingObjective.Override(ctx, &st.SV, int64(LBRebalancingCPU))
+		require.Equal(t,
+			LBRebalancingQueries,
+			ResolveLBRebalancingObjective(ctx, st, gossipStoreDescProvider.GetStores()),
+		)
+		return
+	}
+
 	t.Run("latest version supports all rebalance objectives", func(t *testing.T) {
 		st := cluster.MakeTestingClusterSettings()
 		gossipStoreDescProvider := testMakeProviderNotifier(allPositiveCPUMap)
@@ -168,6 +192,26 @@ func TestRebalanceObjectiveManager(t *testing.T) {
 		), &callbacks
 	}
 
+	// On ARM MacOS and other architectures, grunning isn't supported so
+	// changing the objective to CPU should never work. If this test is  run on
+	// one of these unsupported aarch, test this behavior only.
+	if !grunning.Supported() {
+		st := cluster.MakeTestingClusterSettings()
+		LoadBasedRebalancingObjective.Override(ctx, &st.SV, int64(LBRebalancingQueries))
+		providerNotifier := testMakeProviderNotifier(allPositiveCPUMap)
+		manager, callbacks := makeTestManager(st, providerNotifier)
+
+		require.Equal(t, LBRebalancingQueries, manager.Objective())
+
+		// Changing the objective to CPU should not work since it isn't
+		// supported on this aarch.
+		LoadBasedRebalancingObjective.Override(ctx, &st.SV, int64(LBRebalancingCPU))
+		require.Equal(t, LBRebalancingCPU, manager.Objective())
+		require.Len(t, *callbacks, 0)
+
+		return
+	}
+
 	t.Run("latest version", func(t *testing.T) {
 		st := cluster.MakeTestingClusterSettings()
 		LoadBasedRebalancingObjective.Override(ctx, &st.SV, int64(LBRebalancingCPU))
