From 8115d21eae6d25cccbb9308796cac1b2220b6c77 Mon Sep 17 00:00:00 2001
From: Joel Rebello <jrebello@packet.com>
Date: Mon, 2 Sep 2024 12:34:10 +0200
Subject: [PATCH] worker: setup handler for resumed task

This fixes a bug where the resumed task would have no action step
handler defined, resumed tasks are now planned differently so-as to
ensure the previous plan is not overwritten and the handlers are
available for execution.
---
 internal/worker/task_handler.go      | 35 +++++++++++-
 internal/worker/task_handler_test.go | 80 ++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 2 deletions(-)

diff --git a/internal/worker/task_handler.go b/internal/worker/task_handler.go
index 5f427195..58874b8e 100644
--- a/internal/worker/task_handler.go
+++ b/internal/worker/task_handler.go
@@ -30,7 +30,8 @@ var (
 //
 // The handler is instantiated to run a single task
 type handler struct {
-	mode model.RunMode
+	mode    model.RunMode
+	resumed bool
 	*runner.TaskHandlerContext
 }
 
@@ -42,7 +43,8 @@ func newHandler(
 	logger *logrus.Entry,
 ) runner.TaskHandler {
 	return &handler{
-		mode: mode,
+		mode:    mode,
+		resumed: task.State == model.StateActive,
 		TaskHandlerContext: &runner.TaskHandlerContext{
 			Task:      task,
 			Publisher: publisher,
@@ -141,6 +143,10 @@ func (t handler) inventoryInband(ctx context.Context) (*common.Device, error) {
 }
 
 func (t *handler) PlanActions(ctx context.Context) error {
+	if t.resumed && len(t.Task.Data.ActionsPlanned) > 0 {
+		return t.planResumedTask()
+	}
+
 	switch t.Task.Data.FirmwarePlanMethod {
 	case model.FromFirmwareSet:
 		return t.planFromFirmwareSet(ctx)
@@ -195,6 +201,31 @@ func (t *handler) planFromFirmwareSet(ctx context.Context) error {
 	return nil
 }
 
+func (t *handler) planResumedTask() error {
+	if t.mode == model.RunOutofband {
+		return errors.Wrap(errTaskPlanActions, "resume task not (yet) supported on out-of-band firmware installs")
+	}
+
+	for _, action := range t.Task.Data.ActionsPlanned {
+		if rctypes.StateIsComplete(action.State) {
+			continue
+		}
+
+		actionCtx := &runner.ActionHandlerContext{
+			TaskHandlerContext: t.TaskHandlerContext,
+			Firmware:           &action.Firmware,
+			First:              action.First,
+			Last:               action.Last,
+		}
+
+		if err := inband.AssignStepHandlers(action, actionCtx); err != nil {
+			return errors.Wrap(errTaskPlanActions, "failed to assign action step handler: "+err.Error())
+		}
+	}
+
+	return nil
+}
+
 // planInstall sets up the firmware install plan
 //
 // This returns a list of actions to added to the task and a list of action state machines for those actions.
diff --git a/internal/worker/task_handler_test.go b/internal/worker/task_handler_test.go
index 82c2a29f..eaf697e3 100644
--- a/internal/worker/task_handler_test.go
+++ b/internal/worker/task_handler_test.go
@@ -422,3 +422,83 @@ func TestPlanInstall_Inband(t *testing.T) {
 	require.Equal(t, "drive", actions[0].Firmware.Component, "expect drive component action")
 	require.Equal(t, "nic", actions[1].Firmware.Component, "expect nic component action")
 }
+
+func TestPlanResumedTask(t *testing.T) {
+	t.Parallel()
+
+	logger := logrus.NewEntry(logrus.New())
+
+	serverID := uuid.MustParse("fa125199-e9dd-47d4-8667-ce1d26f58c4a")
+	taskID := uuid.MustParse("05c3296d-be5d-473a-b90c-4ce66cfdec65")
+	taskHandlerCtx := &runner.TaskHandlerContext{
+		Logger: logger,
+		Task: &model.Task{
+			ID:       taskID,
+			WorkerID: registry.GetID("test-app").String(),
+			Data: &model.TaskData{
+				ActionsPlanned: []*model.Action{
+					{
+						Firmware: rctypes.Firmware{Component: "drive", Version: "4.2.1"},
+						Steps: []*model.Step{
+							{
+								Name:  "downloadFirmware",
+								State: model.StateSucceeded,
+							},
+							{
+								Name:  "installFirmware",
+								State: model.StateSucceeded,
+							},
+						},
+					},
+					{
+						Firmware: rctypes.Firmware{Component: "nic", Version: "1.2.2"},
+						Steps: []*model.Step{
+							{
+								Name:  "installFirmware",
+								State: model.StateSucceeded,
+							},
+							{
+								Name:  "powerCycleServer",
+								State: model.StateActive,
+							},
+						},
+					},
+				},
+			},
+			Parameters: &rctypes.FirmwareInstallTaskParameters{
+				AssetID:      serverID,
+				ForceInstall: true,
+			},
+			Server: &rtypes.Server{
+				ID: serverID.String(),
+				Components: rtypes.Components{
+					{
+						Name:     "nic",
+						Model:    "0001",
+						Firmware: &common.Firmware{Installed: "1.2.2"},
+					},
+					{
+						Name:     "drive",
+						Model:    "000",
+						Firmware: &common.Firmware{Installed: "4.2.1"},
+					},
+				},
+			},
+		},
+	}
+
+	h := handler{mode: model.RunInband, TaskHandlerContext: taskHandlerCtx}
+	err := h.planResumedTask()
+	require.NoError(t, err, "no errors returned")
+
+	require.Equal(t, 2, len(taskHandlerCtx.Task.Data.ActionsPlanned), "expect 2 actions to be intact")
+	// expect non-nil handlers for each step
+	for _, action := range taskHandlerCtx.Task.Data.ActionsPlanned {
+		require.Equal(t, 2, len(action.Steps), "expect 2 steps in each action")
+		for _, step := range action.Steps {
+			if step.State == rctypes.Active {
+				require.NotNil(t, step.Handler)
+			}
+		}
+	}
+}