added unit test

Signed-off-by: Daniel Rammer <[email protected]>
flyteorg · Dec 3, 2024 · 22a8337 · 22a8337
1 parent f34fb8f
commit 22a8337
Showing 1 changed file with 66 additions and 19 deletions.
diff --git a/flytepropeller/pkg/controller/nodes/array/handler_test.go b/flytepropeller/pkg/controller/nodes/array/handler_test.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 
 	idlcore "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
@@ -184,9 +186,15 @@ func createNodeExecutionContext(dataStore *storage.DataStore, eventRecorder inte
 	nCtx.OnNodeStateWriter().Return(nodeStateWriter)
 
 	// NodeStatus
+	nowMinus := time.Now().Add(time.Duration(-5) * time.Second)
+	metav1NowMinus := metav1.Time{
+		Time: nowMinus,
+	}
 	nCtx.OnNodeStatus().Return(&v1alpha1.NodeStatus{
 		DataDir:   storage.DataReference("s3://bucket/data"),
 		OutputDir: storage.DataReference("s3://bucket/output"),
+		LastAttemptStartedAt: &metav1NowMinus,
+		StartedAt: &metav1NowMinus,
 	})
 
 	return nCtx
@@ -508,25 +516,27 @@ func TestHandleArrayNodePhaseExecuting(t *testing.T) {
 	}
 
 	tests := []struct {
-		name                           string
-		parallelism                    *uint32
-		minSuccessRatio                *float32
-		subNodePhases                  []v1alpha1.NodePhase
-		subNodeTaskPhases              []core.Phase
-		subNodeTransitions             []handler.Transition
-		expectedArrayNodePhase         v1alpha1.ArrayNodePhase
-		expectedArrayNodeSubPhases     []v1alpha1.NodePhase
-		expectedTransitionPhase        handler.EPhase
-		expectedExternalResourcePhases []idlcore.TaskExecution_Phase
-		currentWfParallelism           uint32
-		maxWfParallelism               uint32
-		incrementParallelismCount      uint32
-		useFakeEventRecorder           bool
-		eventRecorderFailures          uint32
-		eventRecorderError             error
-		expectedTaskPhaseVersion       uint32
-		expectHandleError              bool
-		expectedEventingCalls          int
+		name                                string
+		parallelism                         *uint32
+		minSuccessRatio                     *float32
+		subNodePhases                       []v1alpha1.NodePhase
+		subNodeTaskPhases                   []core.Phase
+		subNodeDeltaTimestamps              []uint64
+		subNodeTransitions                  []handler.Transition
+		expectedArrayNodePhase              v1alpha1.ArrayNodePhase
+		expectedArrayNodeSubPhases          []v1alpha1.NodePhase
+		expectedDiffArrayNodeSubDeltaTimestamps []bool
+		expectedTransitionPhase             handler.EPhase
+		expectedExternalResourcePhases      []idlcore.TaskExecution_Phase
+		currentWfParallelism                uint32
+		maxWfParallelism                    uint32
+		incrementParallelismCount           uint32
+		useFakeEventRecorder                bool
+		eventRecorderFailures               uint32
+		eventRecorderError                  error
+		expectedTaskPhaseVersion            uint32
+		expectHandleError                   bool
+		expectedEventingCalls               int
 	}{
 		{
 			name:        "StartAllSubNodes",
@@ -829,6 +839,31 @@ func TestHandleArrayNodePhaseExecuting(t *testing.T) {
 			expectHandleError:              true,
 			expectedEventingCalls:          1,
 		},
+		{
+			name:        "DeltaTimestampUpdates",
+			parallelism: uint32Ptr(0),
+			subNodePhases: []v1alpha1.NodePhase{
+				v1alpha1.NodePhaseQueued,
+				v1alpha1.NodePhaseRunning,
+			},
+			subNodeTaskPhases: []core.Phase{
+				core.PhaseUndefined,
+				core.PhaseUndefined,
+			},
+			subNodeTransitions: []handler.Transition{
+				handler.DoTransition(handler.TransitionTypeEphemeral, handler.PhaseInfoRunning(&handler.ExecutionInfo{})),
+				handler.DoTransition(handler.TransitionTypeEphemeral, handler.PhaseInfoRetryableFailure(idlcore.ExecutionError_SYSTEM, "", "", &handler.ExecutionInfo{})),
+			},
+			expectedArrayNodePhase: v1alpha1.ArrayNodePhaseExecuting,
+			expectedArrayNodeSubPhases: []v1alpha1.NodePhase{
+				v1alpha1.NodePhaseRunning,
+				v1alpha1.NodePhaseRetryableFailure,
+			},
+			expectedTaskPhaseVersion:       1,
+			expectedTransitionPhase:        handler.EPhaseRunning,
+			expectedExternalResourcePhases: []idlcore.TaskExecution_Phase{idlcore.TaskExecution_RUNNING, idlcore.TaskExecution_FAILED},
+			incrementParallelismCount:      1,
+		},
 	}
 
 	for _, test := range tests {
@@ -870,6 +905,10 @@ func TestHandleArrayNodePhaseExecuting(t *testing.T) {
 				arrayNodeState.SubNodePhases.SetItem(i, bitarray.Item(nodePhase)) // #nosec G115
 			}
 
+			for i, deltaTimestmap := range test.subNodeDeltaTimestamps {
+				arrayNodeState.SubNodeDeltaTimestamps.SetItem(i, bitarray.Item(deltaTimestmap)) // #nosec G115
+			}
+
 			nodeSpec := arrayNodeSpec
 			nodeSpec.ArrayNode.Parallelism = test.parallelism
 			nodeSpec.ArrayNode.MinSuccessRatio = test.minSuccessRatio
@@ -925,6 +964,14 @@ func TestHandleArrayNodePhaseExecuting(t *testing.T) {
 				assert.Equal(t, expectedPhase, v1alpha1.NodePhase(arrayNodeState.SubNodePhases.GetItem(i))) // #nosec G115
 			}
 
+			for i, expectedDiffDeltaTimestamps := range test.expectedDiffArrayNodeSubDeltaTimestamps {
+				if expectedDiffDeltaTimestamps {
+					assert.NotEqual(t, arrayNodeState.SubNodeDeltaTimestamps.GetItem(i), test.subNodeDeltaTimestamps[i])
+				} else {
+					assert.Equal(t, arrayNodeState.SubNodeDeltaTimestamps.GetItem(i), test.subNodeDeltaTimestamps[i])
+				}
+			}
+
 			bufferedEventRecorder, ok := eventRecorder.(*bufferedEventRecorder)
 			if ok {
 				if len(test.expectedExternalResourcePhases) > 0 {