Fix more tests

Tom-Newton · Dec 13, 2023 · 14ecbd9 · 14ecbd9
1 parent 2509e0e
commit 14ecbd9
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/flytepropeller/pkg/controller/nodes/executor.go b/flytepropeller/pkg/controller/nodes/executor.go
@@ -1279,7 +1279,7 @@ func (c *nodeExecutor) handleQueuedOrRunningNode(ctx context.Context, nCtx inter
 func (c *nodeExecutor) handleRetryableFailure(ctx context.Context, nCtx interfaces.NodeExecutionContext, h interfaces.NodeHandler) (interfaces.NodeStatus, error) {
 	nodeStatus := nCtx.NodeStatus()
 	logger.Debugf(ctx, "node failed with retryable failure, aborting and finalizing, message: %s", nodeStatus.GetMessage())
-	if err := c.Abort(ctx, h, nCtx, nodeStatus.GetExecutionError().GetMessage(), false); err != nil {
+	if err := c.Abort(ctx, h, nCtx, nodeStatus.GetMessage(), false); err != nil {
 		return interfaces.NodeStatusUndefined, err
 	}
 

diff --git a/flytepropeller/pkg/controller/nodes/executor_test.go b/flytepropeller/pkg/controller/nodes/executor_test.go
@@ -582,7 +582,8 @@ func TestNodeExecutor_RecursiveNodeHandler_Recurse(t *testing.T) {
 
 			mockN2Status.OnGetStoppedAt().Return(nil)
 			var ee *core.ExecutionError
-			mockN2Status.On("UpdatePhase", expectedN2Phase, mock.Anything, mock.AnythingOfType("string"), ee)
+			// TODO: Add tests case with clearStateOnTermination=true
+			mockN2Status.On("UpdatePhase", expectedN2Phase, mock.Anything, mock.AnythingOfType("string"), false, ee)
 			mockN2Status.OnIsDirty().Return(false)
 			mockN2Status.OnGetTaskNodeStatus().Return(nil)
 			mockN2Status.On("ClearDynamicNodeStatus").Return(nil)
@@ -1329,7 +1330,7 @@ func TestNodeExecutor_RecursiveNodeHandler_BranchNode(t *testing.T) {
 
 				if test.phaseUpdateExpected {
 					var ee *core.ExecutionError
-					branchTakeNodeStatus.On("UpdatePhase", v1alpha1.NodePhaseQueued, mock.Anything, mock.Anything, ee).Return()
+					branchTakeNodeStatus.On("UpdatePhase", v1alpha1.NodePhaseQueued, mock.Anything, mock.Anything, false, ee).Return()
 				}
 
 				leafDag := executors.NewLeafNodeDAGStructure(branchTakenNodeID, parentBranchNodeID)

diff --git a/flytepropeller/pkg/controller/nodes/transformers.go b/flytepropeller/pkg/controller/nodes/transformers.go
@@ -228,7 +228,7 @@ func ToK8sTime(t time.Time) v1.Time {
 	return v1.Time{Time: t}
 }
 
-func UpdateNodeStatus(np v1alpha1.NodePhase, p handler.PhaseInfo, n interfaces.NodeStateReader, s v1alpha1.ExecutableNodeStatus clearStateOnTermination, bool) {
+func UpdateNodeStatus(np v1alpha1.NodePhase, p handler.PhaseInfo, n interfaces.NodeStateReader, s v1alpha1.ExecutableNodeStatus, clearStateOnTermination bool) {
 	// We update the phase and / or reason only if they are not already updated
 	if np != s.GetPhase() || p.GetReason() != s.GetMessage() {
 		s.UpdatePhase(np, ToK8sTime(p.GetOccurredAt()), p.GetReason(), clearStateOnTermination, p.GetErr())