update unit tests

Signed-off-by: Paul Dittamo <pvdittamo@gmail.com>
flyteorg · Jan 13, 2025 · 8e4f4fa · 8e4f4fa
1 parent 6cc769f
commit 8e4f4fa
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 24 deletions.
diff --git a/flytepropeller/pkg/controller/nodes/array/handler.go b/flytepropeller/pkg/controller/nodes/array/handler.go
@@ -466,18 +466,11 @@ func (a *arrayNodeHandler) Handle(ctx context.Context, nCtx interfaces.NodeExecu
 			nCtx.ExecutionContext().IncrementParallelism()
 		}
 	case v1alpha1.ArrayNodePhaseFailing:
+		// note: sub node eventing handled during Abort
 		if err := a.Abort(ctx, nCtx, "ArrayNodeFailing"); err != nil {
 			return handler.UnknownTransition, err
 		}
 
-		// ensure task_execution set to failed - this should already be sent by the abort handler
-		if err := eventRecorder.finalize(ctx, nCtx, idlcore.TaskExecution_FAILED, 0, a.eventConfig); err != nil {
-			if !eventsErr.IsAlreadyExists(err) {
-				logger.Errorf(ctx, "ArrayNode event recording failed: [%s]", err.Error())
-				return handler.UnknownTransition, err
-			}
-		}
-
 		// fail with reported error if one exists
 		if arrayNodeState.Error != nil {
 			return handler.DoTransition(handler.TransitionTypeEphemeral, handler.PhaseInfoFailureErr(arrayNodeState.Error, nil)), nil

diff --git a/flytepropeller/pkg/controller/nodes/array/handler_test.go b/flytepropeller/pkg/controller/nodes/array/handler_test.go
@@ -202,40 +202,56 @@ func createNodeExecutionContext(dataStore *storage.DataStore, eventRecorder inte
 
 func TestAbort(t *testing.T) {
 	ctx := context.Background()
-	scope := promutils.NewTestScope()
-	dataStore, err := storage.NewDataStore(&storage.Config{
-		Type: storage.TypeMemory,
-	}, scope)
-	assert.NoError(t, err)
-
-	nodeHandler := &mocks.NodeHandler{}
-	nodeHandler.OnAbortMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	nodeHandler.OnFinalizeMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil)
-
-	// initialize ArrayNodeHandler
-	arrayNodeHandler, err := createArrayNodeHandler(ctx, t, nodeHandler, dataStore, scope)
-	assert.NoError(t, err)
 
 	tests := []struct {
 		name                           string
 		inputMap                       map[string][]int64
 		subNodePhases                  []v1alpha1.NodePhase
 		subNodeTaskPhases              []core.Phase
 		expectedExternalResourcePhases []idlcore.TaskExecution_Phase
+		arrayNodeState                 v1alpha1.ArrayNodePhase
+		expectedTaskExecutionPhase     idlcore.TaskExecution_Phase
 	}{
 		{
-			name: "Success",
+			name: "Aborted after failed",
+			inputMap: map[string][]int64{
+				"foo": []int64{0, 1, 2},
+			},
+			subNodePhases:                  []v1alpha1.NodePhase{v1alpha1.NodePhaseSucceeded, v1alpha1.NodePhaseRunning, v1alpha1.NodePhaseNotYetStarted},
+			subNodeTaskPhases:              []core.Phase{core.PhaseSuccess, core.PhaseRunning, core.PhaseUndefined},
+			expectedExternalResourcePhases: []idlcore.TaskExecution_Phase{idlcore.TaskExecution_ABORTED},
+			arrayNodeState:                 v1alpha1.ArrayNodePhaseFailing,
+			expectedTaskExecutionPhase:     idlcore.TaskExecution_FAILED,
+		},
+		{
+			name: "Aborted while running",
 			inputMap: map[string][]int64{
 				"foo": []int64{0, 1, 2},
 			},
 			subNodePhases:                  []v1alpha1.NodePhase{v1alpha1.NodePhaseSucceeded, v1alpha1.NodePhaseRunning, v1alpha1.NodePhaseNotYetStarted},
 			subNodeTaskPhases:              []core.Phase{core.PhaseSuccess, core.PhaseRunning, core.PhaseUndefined},
 			expectedExternalResourcePhases: []idlcore.TaskExecution_Phase{idlcore.TaskExecution_ABORTED},
+			arrayNodeState:                 v1alpha1.ArrayNodePhaseExecuting,
+			expectedTaskExecutionPhase:     idlcore.TaskExecution_ABORTED,
 		},
 	}
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
+			scope := promutils.NewTestScope()
+			dataStore, err := storage.NewDataStore(&storage.Config{
+				Type: storage.TypeMemory,
+			}, scope)
+			assert.NoError(t, err)
+
+			nodeHandler := &mocks.NodeHandler{}
+			nodeHandler.OnAbortMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			nodeHandler.OnFinalizeMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+			// initialize ArrayNodeHandler
+			arrayNodeHandler, err := createArrayNodeHandler(ctx, t, nodeHandler, dataStore, scope)
+			assert.NoError(t, err)
+
 			// initialize universal variables
 			literalMap := convertMapToArrayLiterals(test.inputMap)
 
@@ -250,7 +266,7 @@ func TestAbort(t *testing.T) {
 
 			// initialize ArrayNodeState
 			arrayNodeState := &handler.ArrayNodeState{
-				Phase: v1alpha1.ArrayNodePhaseFailing,
+				Phase: test.arrayNodeState,
 			}
 			for _, item := range []struct {
 				arrayReference *bitarray.CompactArray
@@ -279,12 +295,13 @@ func TestAbort(t *testing.T) {
 			nCtx := createNodeExecutionContext(dataStore, eventRecorder, nil, literalMap, &arrayNodeSpec, arrayNodeState, 0, workflowMaxParallelism)
 
 			// evaluate node
-			err := arrayNodeHandler.Abort(ctx, nCtx, "foo")
+			err = arrayNodeHandler.Abort(ctx, nCtx, "foo")
 			assert.NoError(t, err)
 
 			nodeHandler.AssertNumberOfCalls(t, "Abort", len(test.expectedExternalResourcePhases))
 			if len(test.expectedExternalResourcePhases) > 0 {
 				assert.Equal(t, 1, len(eventRecorder.taskExecutionEvents))
+				assert.Equal(t, test.expectedTaskExecutionPhase, eventRecorder.taskExecutionEvents[0].GetPhase())
 
 				externalResources := eventRecorder.taskExecutionEvents[0].GetMetadata().GetExternalResources()
 				assert.Equal(t, len(test.expectedExternalResourcePhases), len(externalResources))
@@ -1296,6 +1313,9 @@ func TestHandleArrayNodePhaseSucceeding(t *testing.T) {
 					assert.Equal(t, int64(*outputValue), collection.GetLiterals()[i].GetScalar().GetPrimitive().GetInteger())
 				}
 			}
+
+			assert.Equal(t, 1, len(eventRecorder.taskExecutionEvents))
+			assert.Equal(t, idlcore.TaskExecution_SUCCEEDED, eventRecorder.taskExecutionEvents[0].GetPhase())
 		})
 	}
 }
@@ -1374,6 +1394,9 @@ func TestHandleArrayNodePhaseFailing(t *testing.T) {
 			assert.Equal(t, test.expectedArrayNodePhase, arrayNodeState.Phase)
 			assert.Equal(t, test.expectedTransitionPhase, transition.Info().GetPhase())
 			nodeHandler.AssertNumberOfCalls(t, "Abort", test.expectedAbortCalls)
+
+			assert.Equal(t, 1, len(eventRecorder.taskExecutionEvents))
+			assert.Equal(t, idlcore.TaskExecution_FAILED, eventRecorder.taskExecutionEvents[0].GetPhase())
 		})
 	}
 }