Fix --clusterPool argument override (flyteorg#360)

* fix clusterPool override * increase coverage Signed-off-by: Iaroslav Ciupin <[email protected]>
austin362667 · May 7, 2024 · 2c39143 · 2c39143
1 parent b2c3bf3
commit 2c39143
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/flytectl/cmd/create/execution_util.go b/flytectl/cmd/create/execution_util.go
@@ -201,6 +201,9 @@ func resolveOverrides(toBeOverridden *ExecutionConfig, project string, domain st
 	if executionConfig.Version != "" {
 		toBeOverridden.Version = executionConfig.Version
 	}
+	if executionConfig.ClusterPool != "" {
+		toBeOverridden.ClusterPool = executionConfig.ClusterPool
+	}
 	// Use the root project and domain to launch the task/workflow if target is unspecified
 	if executionConfig.TargetProject == "" {
 		toBeOverridden.TargetProject = project

diff --git a/flytectl/cmd/create/execution_util_test.go b/flytectl/cmd/create/execution_util_test.go
@@ -5,10 +5,11 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/flyteorg/flytectl/cmd/config"
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/admin"
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
 
+	"github.com/flyteorg/flytectl/cmd/config"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 )
@@ -192,3 +193,22 @@ func TestCreateExecutionRequestForTask(t *testing.T) {
 		executionConfig.KubeServiceAcct = ""
 	})
 }
+
+func Test_resolveOverrides(t *testing.T) {
+	executionConfig.KubeServiceAcct = "k8s-acct"
+	executionConfig.IamRoleARN = "iam-role"
+	executionConfig.TargetProject = "t-proj"
+	executionConfig.TargetDomain = "t-domain"
+	executionConfig.Version = "v1"
+	executionConfig.ClusterPool = "gpu"
+	cfg := &ExecutionConfig{}
+
+	resolveOverrides(cfg, "p1", "d1")
+
+	assert.Equal(t, "k8s-acct", cfg.KubeServiceAcct)
+	assert.Equal(t, "iam-role", cfg.IamRoleARN)
+	assert.Equal(t, "t-proj", cfg.TargetProject)
+	assert.Equal(t, "t-domain", cfg.TargetDomain)
+	assert.Equal(t, "v1", cfg.Version)
+	assert.Equal(t, "gpu", cfg.ClusterPool)
+}