Fixed tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · Feb 7, 2022 · 0dd93c2 · 0dd93c2
1 parent e75e0ed
commit 0dd93c2
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/go/tasks/plugins/array/catalog.go b/go/tasks/plugins/array/catalog.go
@@ -21,7 +21,6 @@ import (
 	idlCore "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
 )
 
-
 // DetermineDiscoverability checks if there are any previously cached tasks. If there are we will only submit an
 // ArrayJob for the non-cached tasks. The ArrayJob is now a different size, and each task will get a new index location
 // which is different than their original location. To find the original index we construct an indexLookup array.

diff --git a/go/tasks/plugins/array/core/state.go b/go/tasks/plugins/array/core/state.go
@@ -143,7 +143,7 @@ const (
 )
 
 func ToArrayJob(taskTemplate *idlCore.TaskTemplate, taskTypeVersion int32) (*ArrayJob, error) {
-	if taskTemplate.GetConfig() != nil {
+	if taskTemplate != nil && taskTemplate.GetConfig() != nil {
 		config := taskTemplate.GetConfig()
 		arrayJob := &ArrayJob{}
 		var err error
@@ -163,7 +163,7 @@ func ToArrayJob(taskTemplate *idlCore.TaskTemplate, taskTypeVersion int32) (*Arr
 	}
 
 	// Keep backward compatibility for those who use arrayJob proto
-	if taskTemplate.GetCustom() != nil {
+	if taskTemplate != nil && taskTemplate.GetCustom() != nil {
 		arrayJob := &idlPlugins.ArrayJob{}
 		err := utils.UnmarshalStruct(taskTemplate.GetCustom(), arrayJob)
 		if err != nil {
@@ -177,7 +177,7 @@ func ToArrayJob(taskTemplate *idlCore.TaskTemplate, taskTypeVersion int32) (*Arr
 		}, nil
 	}
 
-	if taskTypeVersion == 0 || taskTemplate.Type == AwsBatchTaskType {
+	if taskTypeVersion == 0 || (taskTemplate != nil && taskTemplate.Type == AwsBatchTaskType) {
 		return &ArrayJob{
 			Parallelism:  1,
 			Size:         1,

diff --git a/go/tasks/plugins/array/core/state_test.go b/go/tasks/plugins/array/core/state_test.go
@@ -312,6 +312,17 @@ func TestToArrayJob(t *testing.T) {
 		})
 	})
 
+	t.Run("task_type_version == AwsBatchTaskType", func(t *testing.T) {
+		taskTemplate := &idlCore.TaskTemplate{Type: AwsBatchTaskType}
+		arrayJob, err := ToArrayJob(taskTemplate, 1)
+		assert.NoError(t, err)
+		assert.True(t, *arrayJob == ArrayJob{
+			Parallelism:  1,
+			Size:         1,
+			MinSuccesses: 1,
+		})
+	})
+
 	t.Run("ToArrayJob with config", func(t *testing.T) {
 		config := map[string]string{
 			"Parallelism":     "10",