From f14bee845694c580608c1d46f17b56c0f12aa3ed Mon Sep 17 00:00:00 2001
From: Amogh Rathore <amoghr@amazon.com>
Date: Fri, 16 Jun 2023 16:08:33 -0700
Subject: [PATCH] Add tests for Task Protection endpoints for the case when ECS
 returns more than one failure (#3749)

---
 agent/handlers/task_server_setup_test.go | 54 ++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/agent/handlers/task_server_setup_test.go b/agent/handlers/task_server_setup_test.go
index a36a00b16c8..fdad96c409c 100644
--- 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,