fix flaky TestHandleEBSAckTimeout test for windows

aws · Sep 12, 2023 · e8599e1 · e8599e1
1 parent e7e9eca
commit e8599e1
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/agent/ebs/watcher_test.go b/agent/ebs/watcher_test.go
@@ -309,10 +309,17 @@ func TestHandleEBSAckTimeout(t *testing.T) {
 	watcher := newTestEBSWatcher(ctx, taskEngineState, eventChannel, mockDiscoveryClient)
 
 	watcher.HandleResourceAttachment(ebsAttachment)
-	time.Sleep(10 * time.Millisecond)
-	assert.Len(t, taskEngineState.(*dockerstate.DockerTaskEngineState).GetAllEBSAttachments(), 0)
-	ebsAttachment, ok := taskEngineState.(*dockerstate.DockerTaskEngineState).GetEBSByVolumeId(volumeID)
-	assert.False(t, ok)
+
+	for {
+		time.Sleep(time.Millisecond * 5)
+		if len(taskEngineState.(*dockerstate.DockerTaskEngineState).GetAllEBSAttachments()) == 0 {
+			// TODO Include a data client check. This will be introduced in a near future PR
+			_, ok := taskEngineState.(*dockerstate.DockerTaskEngineState).GetEBSByVolumeId(volumeID)
+			assert.False(t, ok)
+			break
+		}
+	}
+
 }
 
 // TestHandleMismatchEBSAttachment tests handling an EBS attachment but found a different volume attached