Adding enableFaultInjection to task payload

aws · Sep 25, 2024 · 19fad6a · 19fad6a
1 parent 08341f7
commit 19fad6a
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 16 deletions.
diff --git a/agent/acs/session/payload_responder.go b/agent/acs/session/payload_responder.go
@@ -132,6 +132,12 @@ func (pmHandler *payloadMessageHandler) addPayloadTasks(payload *ecsacs.PayloadM
 			loggerfield.DesiredStatus: apiTask.GetDesiredStatus(),
 		})
 
+		if apiTask.IsFaultInjectionEnabled() {
+			logger.Info("Fault Injection Enabled for task", logger.Fields{
+				loggerfield.TaskARN: apiTask.Arn,
+			})
+		}
+
 		if task.RoleCredentials != nil {
 			// The payload from ACS for the task has credentials for the
 			// task. Add those to the credentials manager and set the

diff --git a/agent/api/task/task.go b/agent/api/task/task.go
@@ -299,8 +299,7 @@ type Task struct {
 
 	NetworkNamespace string `json:"NetworkNamespace,omitempty"`
 
-	// TODO: Will need to initialize/set the value in a follow PR
-	FaultInjectionEnabled bool `json:"FaultInjectionEnabled,omitempty"`
+	EnableFaultInjection bool `json:"enableFaultInjection,omitempty"`
 
 	// DefaultIfname is used to reference the default network interface name on the task network namespace
 	// For AWSVPC mode, it can be eth0 which corresponds to the interface name on the task ENI
@@ -322,6 +321,9 @@ func TaskFromACS(acsTask *ecsacs.Task, envelope *ecsacs.PayloadMessage) (*Task,
 		return nil, err
 	}
 
+	// Set the EnableFaultInjection field if present
+	task.EnableFaultInjection = aws.BoolValue(acsTask.EnableFaultInjection)
+
 	// Overrides the container command if it's set
 	for _, container := range task.Containers {
 		if (container.Overrides != apicontainer.ContainerOverrides{}) && container.Overrides.Command != nil {
@@ -3771,7 +3773,7 @@ func (task *Task) IsFaultInjectionEnabled() bool {
 	task.lock.RLock()
 	defer task.lock.RUnlock()
 
-	return task.FaultInjectionEnabled
+	return task.EnableFaultInjection
 }
 
 func (task *Task) GetNetworkMode() string {

diff --git a/agent/handlers/task_server_setup_test.go b/agent/handlers/task_server_setup_test.go
@@ -456,9 +456,9 @@ var (
 		}},
 	})
 
-	agentStateExpectations = func(state *mock_dockerstate.MockTaskEngineState, faultInjectionEnabled bool, networkMode string) {
+	agentStateExpectations = func(state *mock_dockerstate.MockTaskEngineState, enableFaultInjection bool, networkMode string) {
 		task := standardTask()
-		task.FaultInjectionEnabled = faultInjectionEnabled
+		task.EnableFaultInjection = enableFaultInjection
 		task.NetworkMode = networkMode
 		task.NetworkNamespace = networkNamespace
 		task.DefaultIfname = defaultIfname
@@ -589,7 +589,7 @@ func expectedV4TaskResponse() v4.TaskResponse {
 	)
 }
 
-func expectedV4TaskNetworkConfig(faultInjectionEnabled bool, networkMode, path, deviceName string) *v4.TaskNetworkConfig {
+func expectedV4TaskNetworkConfig(enableFaultInjection bool, networkMode, path, deviceName string) *v4.TaskNetworkConfig {
 	return v4.NewTaskNetworkConfig(networkMode, path, deviceName)
 }
 
@@ -2092,7 +2092,7 @@ func TestV4TaskMetadata(t *testing.T) {
 		testTMDSRequest(t, TMDSTestCase[v4.TaskResponse]{
 			path: v4BasePath + v3EndpointID + "/task",
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
-				task.FaultInjectionEnabled = true
+				task.EnableFaultInjection = true
 				task.NetworkNamespace = networkNamespace
 				task.DefaultIfname = defaultIfname
 				gomock.InOrder(
@@ -2115,7 +2115,7 @@ func TestV4TaskMetadata(t *testing.T) {
 			path: v4BasePath + v3EndpointID + "/task",
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				hostTask := standardHostTask()
-				hostTask.FaultInjectionEnabled = true
+				hostTask.EnableFaultInjection = true
 				hostTask.NetworkNamespace = networkNamespace
 				hostTask.DefaultIfname = defaultIfname
 				gomock.InOrder(
@@ -3732,9 +3732,9 @@ type networkFaultTestCase struct {
 	expectedStatusCode    int
 	requestBody           interface{}
 	expectedFaultResponse faulttype.NetworkFaultInjectionResponse
-	setStateExpectations  func(state *mock_dockerstate.MockTaskEngineState, faultInjectionEnabled bool, networkMode string)
+	setStateExpectations  func(state *mock_dockerstate.MockTaskEngineState, enableFaultInjection bool, networkMode string)
 	setExecExpectations   execExpectations
-	faultInjectionEnabled bool
+	enableFaultInjection  bool
 	networkMode           string
 }
 
@@ -3749,7 +3749,7 @@ func generateCommonNetworkFaultInjectionTestCases(requestType, successResponse s
 			expectedFaultResponse: faulttype.NewNetworkFaultInjectionSuccessResponse(successResponse),
 			setStateExpectations:  agentStateExpectations,
 			setExecExpectations:   exec,
-			faultInjectionEnabled: true,
+			enableFaultInjection:  true,
 			networkMode:           apitask.HostNetworkMode,
 		},
 		{
@@ -3759,7 +3759,7 @@ func generateCommonNetworkFaultInjectionTestCases(requestType, successResponse s
 			expectedFaultResponse: faulttype.NewNetworkFaultInjectionSuccessResponse(successResponse),
 			setStateExpectations:  agentStateExpectations,
 			setExecExpectations:   exec,
-			faultInjectionEnabled: true,
+			enableFaultInjection:  true,
 			networkMode:           apitask.AWSVPCNetworkMode,
 		},
 	}
@@ -3909,7 +3909,7 @@ func testRegisterFaultHandler(t *testing.T, tcs []networkFaultTestCase, method,
 			execWrapper := mock_execwrapper.NewMockExec(ctrl)
 
 			if tc.setStateExpectations != nil {
-				tc.setStateExpectations(state, tc.faultInjectionEnabled, tc.networkMode)
+				tc.setStateExpectations(state, tc.enableFaultInjection, tc.networkMode)
 			}
 
 			if tc.setExecExpectations != nil {
@@ -3954,7 +3954,7 @@ func TestV4GetTaskMetadataWithTaskNetworkConfig(t *testing.T) {
 			name: "happy case with awsvpc mode",
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				task := standardTask()
-				task.FaultInjectionEnabled = true
+				task.EnableFaultInjection = true
 				task.NetworkNamespace = networkNamespace
 				task.DefaultIfname = defaultIfname
 				gomock.InOrder(
@@ -3973,7 +3973,7 @@ func TestV4GetTaskMetadataWithTaskNetworkConfig(t *testing.T) {
 			name: "happy case with host mode",
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				hostTask := standardHostTask()
-				hostTask.FaultInjectionEnabled = true
+				hostTask.EnableFaultInjection = true
 				hostTask.NetworkNamespace = networkNamespace
 				hostTask.DefaultIfname = defaultIfname
 				gomock.InOrder(

diff --git a/ecs-agent/tmds/handlers/v4/handlers_test.go b/ecs-agent/tmds/handlers/v4/handlers_test.go
@@ -301,7 +301,7 @@ func TestTaskMetadata(t *testing.T) {
 		expectedTaskResponse := taskResponseWithFaultInjectionEnabled()
 		expectedTaskResponse.CredentialsID = ""            // credentials ID not expected
 		expectedTaskResponse.TaskNetworkConfig = nil       // TaskNetworkConfig is not expected and would be used internally
-		expectedTaskResponse.FaultInjectionEnabled = false // FaultInjectionEnabled is not expected and would be used internally
+		expectedTaskResponse.FaultInjectionEnabled = false // EnableFaultInjection is not expected and would be used internally
 
 		handler, _, agentState, _ := setup(t)
 		agentState.EXPECT().