Fix the deployment model method

pkg/app/piped/controller/scheduler.go

@@ -419,7 +419,7 @@ func (s *scheduler) executeStage(sig executor.StopSignal, ps model.PipelineStage
 	defer func() {
 		// When the piped has been terminated (PS kill) while the stage is still running
 		// we should not mark the log persister as completed.
-		if !model.IsCompletedStage(finalStatus) && sig.Terminated() {
+		if !finalStatus.IsCompleted() && sig.Terminated() {
 			return
 		}
 		lp.Complete(time.Minute)

pkg/app/server/grpcapi/web_api.go

@@ -1025,11 +1025,11 @@ func (a *WebAPI) ApproveStage(ctx context.Context, req *webservice.ApproveStageR
 	if err := a.validateDeploymentBelongsToProject(ctx, req.DeploymentId, claims.Role.ProjectId); err != nil {
 		return nil, err
 	}
-	stage, ok := deployment.StageStatusMap()[req.StageId]
+	stage, ok := deployment.StageMap()[req.StageId]
 	if !ok {
 		return nil, status.Error(codes.FailedPrecondition, "The stage was not found in the deployment")
 	}
-	if model.IsCompletedStage(stage) {
+	if stage.Status.IsCompleted() {
 		return nil, status.Errorf(codes.FailedPrecondition, "Could not approve the stage because it was already completed")
 	}
 

pkg/model/deployment.go

@@ -53,21 +53,12 @@ func IsSuccessfullyCompletedDeployment(status DeploymentStatus) bool {
 	return status == DeploymentStatus_DEPLOYMENT_SUCCESS
 }
 
-// IsCompletedStage checks whether the stage is at a completion state.
-func IsCompletedStage(status StageStatus) bool {
-	if status.String() == "" {
-		return false
-	}
-
-	switch status {
-	case StageStatus_STAGE_SUCCESS:
-		return true
-	case StageStatus_STAGE_FAILURE:
-		return true
-	case StageStatus_STAGE_CANCELLED:
-		return true
-	}
-	return false
+// IsCompleted checks whether the stage is at a completion state.
+func (s StageStatus) IsCompleted() bool {
+	return s == StageStatus_STAGE_SUCCESS ||
+		s == StageStatus_STAGE_FAILURE ||
+		s == StageStatus_STAGE_CANCELLED ||
+		s == StageStatus_STAGE_SKIPPED
 }
 
 // CanUpdateDeploymentStatus checks whether the deployment can transit to the given status.
@@ -113,13 +104,13 @@ func CanUpdateStageStatus(cur, next StageStatus) bool {
 	return false
 }
 
-// StageStatusMap returns the map from id to status of all stages.
-func (d *Deployment) StageStatusMap() map[string]StageStatus {
-	statuses := make(map[string]StageStatus, len(d.Stages))
+// StageMap returns the map of id and the stage.
+func (d *Deployment) StageMap() map[string]*PipelineStage {
+	stage := make(map[string]*PipelineStage, len(d.Stages))
 	for _, s := range d.Stages {
-		statuses[s.Id] = s.Status
+		stage[s.Id] = s
 	}
-	return statuses
+	return stage
 }
 
 // CommitHash returns the hash value of trigger commit.

pkg/model/deployment_test.go

@@ -69,3 +69,44 @@ func TestDeployment_ContainTags(t *testing.T) {
 		})
 	}
 }
+
+func TestDeployment_StageMap(t *testing.T) {
+	testcases := []struct {
+		name       string
+		deployment *Deployment
+		want       map[string]*PipelineStage
+	}{
+		{
+			name: "ok",
+			deployment: &Deployment{
+				Stages: []*PipelineStage{
+					{
+						Id: "stage1",
+					},
+					{
+						Id: "stage2",
+					},
+				},
+			},
+			want: map[string]*PipelineStage{
+				"stage1": {
+					Id: "stage1",
+				},
+				"stage2": {
+					Id: "stage2",
+				},
+			},
+		},
+		{
+			name:       "no stages",
+			deployment: &Deployment{},
+			want:       map[string]*PipelineStage{},
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			got := tc.deployment.StageMap()
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}