Add tests for Task Protection endpoints for the case when ECS returns…

… more than one failure (#3749)
aws · Jun 16, 2023 · f14bee8 · f14bee8
1 parent 7565a9a
commit f14bee8
Showing 1 changed file with 54 additions and 0 deletions.
diff --git a/agent/handlers/task_server_setup_test.go b/agent/handlers/task_server_setup_test.go
@@ -2973,6 +2973,34 @@ func TestGetTaskProtection(t *testing.T) {
 			},
 		})
 	})
+	t.Run("more than one ecs failure", func(t *testing.T) {
+		testTMDSRequest(t, TMDSTestCase[agentapi.TaskProtectionResponse]{
+			path:                              path,
+			setStateExpectations:              happyStateExpectations,
+			setCredentialsManagerExpectations: happyCredentialsManagerExpectations,
+			setTaskProtectionClientFactoryExpectations: taskProtectionClientFactoryExpectations(
+				&ecs.GetTaskProtectionOutput{
+					Failures: []*ecs.Failure{
+						{
+							Arn:    aws.String(taskARN),
+							Reason: aws.String("ecs failure 1"),
+						},
+						{
+							Arn:    aws.String(taskARN),
+							Reason: aws.String("ecs failure 2"),
+						},
+					},
+				}, nil),
+			expectedStatusCode: http.StatusInternalServerError,
+			expectedResponseBody: agentapi.TaskProtectionResponse{
+				Error: &agentapi.ErrorResponse{
+					Arn:     taskARN,
+					Code:    ecs.ErrCodeServerException,
+					Message: "Unexpected error occurred",
+				},
+			},
+		})
+	})
 	t.Run("happy case", func(t *testing.T) {
 		testTMDSRequest(t, TMDSTestCase[agentapi.TaskProtectionResponse]{
 			path:                              path,
@@ -3204,6 +3232,32 @@ func TestUpdateTaskProtection(t *testing.T) {
 			},
 		},
 	}))
+	t.Run("more than on ecs failure", runTest(t, TMDSTestCase[agentapi.TaskProtectionResponse]{
+		requestBody:                       happyReqBody,
+		setStateExpectations:              happyStateExpectations,
+		setCredentialsManagerExpectations: happyCredentialsManagerExpectations,
+		setTaskProtectionClientFactoryExpectations: taskProtectionClientFactoryExpectations(
+			&ecs.UpdateTaskProtectionOutput{
+				Failures: []*ecs.Failure{
+					{
+						Arn:    aws.String(taskARN),
+						Reason: aws.String("ecs failure 1"),
+					},
+					{
+						Arn:    aws.String(taskARN),
+						Reason: aws.String("ecs failure 2"),
+					},
+				},
+			}, nil),
+		expectedStatusCode: http.StatusInternalServerError,
+		expectedResponseBody: agentapi.TaskProtectionResponse{
+			Error: &agentapi.ErrorResponse{
+				Arn:     taskARN,
+				Code:    ecs.ErrCodeServerException,
+				Message: "Unexpected error occurred",
+			},
+		},
+	}))
 	t.Run("empty request", runTest(t, TMDSTestCase[agentapi.TaskProtectionResponse]{
 		requestBody:          map[string]string{},
 		setStateExpectations: happyStateExpectations,