Fix flakey TaskStopVerificationACKResponder integration test (#4292)

aws · Aug 19, 2024 · 6e5e459 · 6e5e459
1 parent 19e0acf
commit 6e5e459
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 23 deletions.
diff --git a/agent/acs/session/task_stop_verification_ack_responder_integ_test.go b/agent/acs/session/task_stop_verification_ack_responder_integ_test.go
@@ -82,16 +82,19 @@ func TestTaskStopVerificationACKResponder_StopsSpecificTasks(t *testing.T) {
 	taskEngine, done, dockerClient, _ := engine.SetupIntegTestTaskEngine(engine.DefaultTestConfigIntegTest(), nil, t)
 	defer done()
 
+	testEvents := engine.InitTestEventCollection(taskEngine)
+
 	var tasks []*apitask.Task
 	for i := 0; i < 3; i++ {
 		task := engine.CreateTestTask(fmt.Sprintf("test_task_%d", i))
 		createLongRunningContainers(task, 1)
 		go taskEngine.AddTask(task)
 
-		engine.VerifyContainerManifestPulledStateChange(t, taskEngine)
-		engine.VerifyTaskManifestPulledStateChange(t, taskEngine)
-		engine.VerifyContainerRunningStateChange(t, taskEngine)
-		engine.VerifyTaskRunningStateChange(t, taskEngine)
+		containerName := task.Arn + ":" + task.Containers[0].Name
+		engine.VerifyContainerStatus(apicontainerstatus.ContainerManifestPulled, containerName, testEvents, t)
+		engine.VerifyTaskStatus(apitaskstatus.TaskManifestPulled, task.Arn, testEvents, t)
+		engine.VerifyContainerStatus(apicontainerstatus.ContainerRunning, containerName, testEvents, t)
+		engine.VerifyTaskStatus(apitaskstatus.TaskRunning, task.Arn, testEvents, t)
 		tasks = append(tasks, task)
 	}
 
@@ -113,9 +116,10 @@ func TestTaskStopVerificationACKResponder_StopsSpecificTasks(t *testing.T) {
 	})
 
 	// Wait for all state changes before verifying container and task statuses.
-	for i := 0; i < 2; i++ {
-		engine.VerifyContainerStoppedStateChange(t, taskEngine)
-		engine.VerifyTaskStoppedStateChange(t, taskEngine)
+	for _, task := range tasks[1:] {
+		containerName := task.Arn + ":" + task.Containers[0].Name
+		engine.VerifyContainerStatus(apicontainerstatus.ContainerStopped, containerName, testEvents, t)
+		engine.VerifyTaskStatus(apitaskstatus.TaskStopped, task.Arn, testEvents, t)
 	}
 
 	// Verify that the last 2 tasks and their containers have stopped.

diff --git a/agent/engine/common_integ_testutil.go b/agent/engine/common_integ_testutil.go
@@ -355,10 +355,10 @@ func InitTestEventCollection(taskEngine TaskEngine) *TestEvents {
 
 // This method queries the TestEvents struct to check a Task Status.
 // This method will block if there are no more stateChangeEvents from the DockerTaskEngine but is expected
-func VerifyTaskStatus(status apitaskstatus.TaskStatus, taskARN string, testEvents *TestEvents, t *testing.T) error {
+func VerifyTaskStatus(status apitaskstatus.TaskStatus, taskARN string, testEvents *TestEvents, t *testing.T) {
 	for {
 		if _, found := testEvents.RecordedEvents[statechange.TaskEvent][status.String()][taskARN]; found {
-			return nil
+			return
 		}
 		event := <-testEvents.StateChangeEvents
 		RecordTestEvent(testEvents, event)
@@ -367,10 +367,10 @@ func VerifyTaskStatus(status apitaskstatus.TaskStatus, taskARN string, testEvent
 
 // This method queries the TestEvents struct to check a Task Status.
 // This method will block if there are no more stateChangeEvents from the DockerTaskEngine but is expected
-func VerifyContainerStatus(status apicontainerstatus.ContainerStatus, ARNcontName string, testEvents *TestEvents, t *testing.T) error {
+func VerifyContainerStatus(status apicontainerstatus.ContainerStatus, ARNcontName string, testEvents *TestEvents, t *testing.T) {
 	for {
 		if _, found := testEvents.RecordedEvents[statechange.ContainerEvent][status.String()][ARNcontName]; found {
-			return nil
+			return
 		}
 		event := <-testEvents.StateChangeEvents
 		RecordTestEvent(testEvents, event)

diff --git a/agent/engine/engine_sudo_linux_integ_test.go b/agent/engine/engine_sudo_linux_integ_test.go
@@ -228,28 +228,22 @@ func TestFirelensFluentbit(t *testing.T) {
 	testEvents := InitTestEventCollection(taskEngine)
 
 	//Verify logsender container is running
-	err = VerifyContainerStatus(apicontainerstatus.ContainerRunning, testTask.Arn+":logsender", testEvents, t)
-	assert.NoError(t, err, "Verify logsender container is running")
+	VerifyContainerStatus(apicontainerstatus.ContainerRunning, testTask.Arn+":logsender", testEvents, t)
 
 	//Verify firelens container is running
-	err = VerifyContainerStatus(apicontainerstatus.ContainerRunning, testTask.Arn+":firelens", testEvents, t)
-	assert.NoError(t, err, "Verify firelens container is running")
+	VerifyContainerStatus(apicontainerstatus.ContainerRunning, testTask.Arn+":firelens", testEvents, t)
 
 	//Verify task is in running state
-	err = VerifyTaskStatus(apitaskstatus.TaskRunning, testTask.Arn, testEvents, t)
-	assert.NoError(t, err, "Not verified task running")
+	VerifyTaskStatus(apitaskstatus.TaskRunning, testTask.Arn, testEvents, t)
 
 	//Verify logsender container is stopped
-	err = VerifyContainerStatus(apicontainerstatus.ContainerStopped, testTask.Arn+":logsender", testEvents, t)
-	assert.NoError(t, err)
+	VerifyContainerStatus(apicontainerstatus.ContainerStopped, testTask.Arn+":logsender", testEvents, t)
 
 	//Verify firelens container is stopped
-	err = VerifyContainerStatus(apicontainerstatus.ContainerStopped, testTask.Arn+":firelens", testEvents, t)
-	assert.NoError(t, err)
+	VerifyContainerStatus(apicontainerstatus.ContainerStopped, testTask.Arn+":firelens", testEvents, t)
 
 	//Verify the task itself has stopped
-	err = VerifyTaskStatus(apitaskstatus.TaskStopped, testTask.Arn, testEvents, t)
-	assert.NoError(t, err)
+	VerifyTaskStatus(apitaskstatus.TaskStopped, testTask.Arn, testEvents, t)
 
 	taskID := testTask.GetID()