From f5deb20a2dd98c5d2fcea6f1ea149d4fe84c5aed Mon Sep 17 00:00:00 2001
From: Aaron Wilson <aawilson@nvidia.com>
Date: Tue, 20 Aug 2024 10:32:40 -0500
Subject: [PATCH] Go API: Add EnableRebalance and DisableRebalance API
 functions

Signed-off-by: Aaron Wilson <aawilson@nvidia.com>
---
 ais/test/config_test.go | 23 +++++++++++++++++++++++
 api/cluster.go          | 17 +++++++++++++++++
 tools/client.go         | 14 ++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/ais/test/config_test.go b/ais/test/config_test.go
index a0e5cc4bf3..d3e62f194f 100644
--- a/ais/test/config_test.go
+++ b/ais/test/config_test.go
@@ -358,3 +358,26 @@ func TestConfigOverrideAndResetCluster(t *testing.T) {
 			errWMConfigNotExpected, config.Disk.DiskUtilLowWM, daemonConfig.Disk.DiskUtilLowWM)
 	}
 }
+
+func TestConfigRebalance(t *testing.T) {
+	var (
+		oRebalance = tools.GetClusterConfig(t).Rebalance.Enabled
+	)
+	defer tools.SetClusterConfig(t, cos.StrKVs{
+		"rebalance.enabled": strconv.FormatBool(oRebalance),
+	})
+
+	tools.EnableRebalance(t)
+
+	nRebalance := tools.GetClusterConfig(t).Rebalance.Enabled
+
+	if !nRebalance {
+		t.Errorf("Rebalance was not enabled: current value %v, should be: %v", nRebalance, true)
+	}
+	tools.DisableRebalance(t)
+
+	nRebalance = tools.GetClusterConfig(t).Rebalance.Enabled
+	if nRebalance {
+		t.Errorf("Rebalance was not disabled: current value %v, should be: %v", nRebalance, false)
+	}
+}
diff --git a/api/cluster.go b/api/cluster.go
index 32add950a7..63aea2bcbd 100644
--- a/api/cluster.go
+++ b/api/cluster.go
@@ -262,6 +262,23 @@ func SetClusterConfigUsingMsg(bp BaseParams, configToUpdate *cmn.ConfigToSet, tr
 	return err
 }
 
+func setRebalance(bp BaseParams, enabled bool) error {
+	configToSet := &cmn.ConfigToSet{
+		Rebalance: &cmn.RebalanceConfToSet{
+			Enabled: apc.Ptr(enabled),
+		},
+	}
+	return SetClusterConfigUsingMsg(bp, configToSet, false /*transient*/)
+}
+
+func EnableRebalance(bp BaseParams) error {
+	return setRebalance(bp, true)
+}
+
+func DisableRebalance(bp BaseParams) error {
+	return setRebalance(bp, false)
+}
+
 // all nodes: reset configuration to cluster defaults
 func ResetClusterConfig(bp BaseParams) error {
 	return _putCluster(bp, apc.ActMsg{Action: apc.ActResetConfig})
diff --git a/tools/client.go b/tools/client.go
index c943349565..00f199297c 100644
--- a/tools/client.go
+++ b/tools/client.go
@@ -659,6 +659,20 @@ func SetClusterConfigUsingMsg(t *testing.T, toUpdate *cmn.ConfigToSet) {
 	tassert.CheckFatal(t, err)
 }
 
+func EnableRebalance(t *testing.T) {
+	proxyURL := GetPrimaryURL()
+	bp := BaseAPIParams(proxyURL)
+	err := api.EnableRebalance(bp)
+	tassert.CheckError(t, err)
+}
+
+func DisableRebalance(t *testing.T) {
+	proxyURL := GetPrimaryURL()
+	bp := BaseAPIParams(proxyURL)
+	err := api.DisableRebalance(bp)
+	tassert.CheckError(t, err)
+}
+
 func SetRemAisConfig(t *testing.T, nvs cos.StrKVs) {
 	remoteBP := BaseAPIParams(RemoteCluster.URL)
 	err := api.SetClusterConfig(remoteBP, nvs, false /*transient*/)