Merge branch 'master' into feature/open-telemetry

flyteplugins/go/tasks/plugins/webapi/athena/plugin.go

@@ -3,6 +3,7 @@ package athena
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	awsSdk "github.com/aws/aws-sdk-go-v2/aws"
@@ -177,12 +178,19 @@ func (p Plugin) Status(ctx context.Context, tCtx webapi.StatusContext) (phase co
 
 func createTaskInfo(queryID string, cfg awsSdk.Config) *core.TaskInfo {
 	timeNow := time.Now()
+	var consoleURL string
+	if strings.Contains(cfg.Region, "gov") {
+		consoleURL = "console.amazonaws-us-gov.com"
+	} else {
+		consoleURL = "console.aws.amazon.com"
+	}
 	return &core.TaskInfo{
 		OccurredAt: &timeNow,
 		Logs: []*idlCore.TaskLog{
 			{
-				Uri: fmt.Sprintf("https://%v.console.aws.amazon.com/athena/home?force&region=%v#query/history/%v",
+				Uri: fmt.Sprintf("https://%v.%v/athena/home?force&region=%v#query/history/%v",
 					cfg.Region,
+					consoleURL,
 					cfg.Region,
 					queryID),
 				Name: "Athena Query Console",

flyteplugins/go/tasks/plugins/webapi/athena/plugin_test.go

@@ -22,3 +22,18 @@ func TestCreateTaskInfo(t *testing.T) {
 	assert.Len(t, taskInfo.ExternalResources, 1)
 	assert.Equal(t, taskInfo.ExternalResources[0].ExternalID, "query_id")
 }
+
+
+func TestCreateTaskInfoGovAWS(t *testing.T) {
+	taskInfo := createTaskInfo("query_id", awsSdk.Config{
+		Region: "us-gov-east-1",
+	})
+	assert.EqualValues(t, []*idlCore.TaskLog{
+		{
+			Uri:  "https://us-gov-east-1.console.amazonaws-us-gov.com/athena/home?force&region=us-gov-east-1#query/history/query_id",
+			Name: "Athena Query Console",
+		},
+	}, taskInfo.Logs)
+	assert.Len(t, taskInfo.ExternalResources, 1)
+	assert.Equal(t, taskInfo.ExternalResources[0].ExternalID, "query_id")
+}

flytepropeller/pkg/apis/flyteworkflow/v1alpha1/array_test.go

@@ -0,0 +1,49 @@
+package v1alpha1
+
+import (
+	"testing"
+)
+
+func TestArrayNodeSpec_GetSubNodeSpec(t *testing.T) {
+	nodeSpec := &NodeSpec{}
+	arrayNodeSpec := ArrayNodeSpec{
+		SubNodeSpec: nodeSpec,
+	}
+
+	if arrayNodeSpec.GetSubNodeSpec() != nodeSpec {
+		t.Errorf("Expected nodeSpec, but got a different value")
+	}
+}
+
+func TestArrayNodeSpec_GetParallelism(t *testing.T) {
+	parallelism := uint32(5)
+	arrayNodeSpec := ArrayNodeSpec{
+		Parallelism: parallelism,
+	}
+
+	if arrayNodeSpec.GetParallelism() != parallelism {
+		t.Errorf("Expected %d, but got %d", parallelism, arrayNodeSpec.GetParallelism())
+	}
+}
+
+func TestArrayNodeSpec_GetMinSuccesses(t *testing.T) {
+	minSuccesses := uint32(3)
+	arrayNodeSpec := ArrayNodeSpec{
+		MinSuccesses: &minSuccesses,
+	}
+
+	if *arrayNodeSpec.GetMinSuccesses() != minSuccesses {
+		t.Errorf("Expected %d, but got %d", minSuccesses, *arrayNodeSpec.GetMinSuccesses())
+	}
+}
+
+func TestArrayNodeSpec_GetMinSuccessRatio(t *testing.T) {
+	minSuccessRatio := float32(0.8)
+	arrayNodeSpec := ArrayNodeSpec{
+		MinSuccessRatio: &minSuccessRatio,
+	}
+
+	if *arrayNodeSpec.GetMinSuccessRatio() != minSuccessRatio {
+		t.Errorf("Expected %f, but got %f", minSuccessRatio, *arrayNodeSpec.GetMinSuccessRatio())
+	}
+}

flytepropeller/pkg/apis/flyteworkflow/v1alpha1/branch_test.go

@@ -1,20 +1,21 @@
-package v1alpha1_test
+package v1alpha1
 
 import (
+	"bytes"
 	"encoding/json"
 	"io/ioutil"
 	"testing"
 
+	"github.com/golang/protobuf/jsonpb"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
-	"github.com/flyteorg/flyte/flytepropeller/pkg/apis/flyteworkflow/v1alpha1"
 )
 
 func TestMarshalUnMarshal_BranchTask(t *testing.T) {
 	r, err := ioutil.ReadFile("testdata/branch.json")
 	assert.NoError(t, err)
-	o := v1alpha1.NodeSpec{}
+	o := NodeSpec{}
 	err = json.Unmarshal(r, &o)
 	assert.NoError(t, err)
 	assert.NotNil(t, o.BranchNode.If)
@@ -25,3 +26,74 @@ func TestMarshalUnMarshal_BranchTask(t *testing.T) {
 		assert.NotEmpty(t, raw)
 	}
 }
+
+// TestBranchNodeSpecMethods tests the methods of the BranchNodeSpec struct.
+func TestErrorMarshalAndUnmarshalJSON(t *testing.T) {
+	coreError := &core.Error{
+		FailedNodeId: "TestNode",
+		Message:      "Test error message",
+	}
+
+	err := Error{Error: coreError}
+	data, jErr := err.MarshalJSON()
+	assert.Nil(t, jErr)
+
+	// Unmarshalling the JSON back to a new core.Error struct
+	var newCoreError core.Error
+	uErr := jsonpb.Unmarshal(bytes.NewReader(data), &newCoreError)
+	assert.Nil(t, uErr)
+	assert.Equal(t, coreError.Message, newCoreError.Message)
+	assert.Equal(t, coreError.FailedNodeId, newCoreError.FailedNodeId)
+}
+
+func TestBranchNodeSpecMethods(t *testing.T) {
+	// Creating a core.BooleanExpression instance for testing
+	boolExpr := &core.BooleanExpression{}
+
+	// Creating an Error instance for testing
+	errorMessage := &core.Error{
+		Message: "Test error",
+	}
+
+	ifNode := NodeID("ifNode")
+	elifNode := NodeID("elifNode")
+	elseNode := NodeID("elseNode")
+
+	// Creating a BranchNodeSpec instance for testing
+	branchNodeSpec := BranchNodeSpec{
+		If: IfBlock{
+			Condition: BooleanExpression{
+				BooleanExpression: boolExpr,
+			},
+			ThenNode: &ifNode,
+		},
+		ElseIf: []*IfBlock{
+			{
+				Condition: BooleanExpression{
+					BooleanExpression: boolExpr,
+				},
+				ThenNode: &elifNode,
+			},
+		},
+		Else:     &elseNode,
+		ElseFail: &Error{Error: errorMessage},
+	}
+
+	assert.Equal(t, boolExpr, branchNodeSpec.If.GetCondition())
+
+	assert.Equal(t, &ifNode, branchNodeSpec.If.GetThenNode())
+
+	assert.Equal(t, &branchNodeSpec.If, branchNodeSpec.GetIf())
+
+	assert.Equal(t, &elseNode, branchNodeSpec.GetElse())
+
+	elifs := branchNodeSpec.GetElseIf()
+	assert.Equal(t, 1, len(elifs))
+	assert.Equal(t, boolExpr, elifs[0].GetCondition())
+	assert.Equal(t, &elifNode, elifs[0].GetThenNode())
+
+	assert.Equal(t, errorMessage, branchNodeSpec.GetElseFail())
+
+	branchNodeSpec.ElseFail = nil
+	assert.Nil(t, branchNodeSpec.GetElseFail())
+}

flytepropeller/pkg/apis/flyteworkflow/v1alpha1/error_test.go

@@ -0,0 +1,30 @@
+package v1alpha1
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
+)
+
+func TestExecutionErrorJSONMarshalling(t *testing.T) {
+	execError := &core.ExecutionError{
+		Code:     "TestCode",
+		Message:  "Test error message",
+		ErrorUri: "Test error uri",
+	}
+
+	execErr := &ExecutionError{ExecutionError: execError}
+	data, jErr := json.Marshal(execErr)
+	assert.Nil(t, jErr)
+
+	newExecErr := &ExecutionError{}
+	uErr := json.Unmarshal(data, newExecErr)
+	assert.Nil(t, uErr)
+
+	assert.Equal(t, execError.Code, newExecErr.ExecutionError.Code)
+	assert.Equal(t, execError.Message, newExecErr.ExecutionError.Message)
+	assert.Equal(t, execError.ErrorUri, newExecErr.ExecutionError.ErrorUri)
+}