wip unit test container restarts

agent/engine/task_manager_test.go

@@ -46,6 +46,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/taskresource/volume"
 	taskresourcevolume "github.com/aws/amazon-ecs-agent/agent/taskresource/volume"
 	apiresource "github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/container/restart"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 	apierrors "github.com/aws/amazon-ecs-agent/ecs-agent/api/errors"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
@@ -1855,6 +1856,131 @@ func TestHandleContainerChangeUpdateMetadataRedundant(t *testing.T) {
 	assert.Equal(t, timeNow, containerCreateTime)
 }
 
+func TestHandleContainerChangeStopped(t *testing.T) {
+	eventStreamName := "TestHandleContainerChangeStopped"
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	containerChangeEventStream := eventstream.NewEventStream(eventStreamName, ctx)
+	containerChangeEventStream.StartListening()
+
+	hostResourceManager := NewHostResourceManager(getTestHostResources())
+	mTask := &managedTask{
+		Task:                       testdata.LoadTask("sleep5"),
+		containerChangeEventStream: containerChangeEventStream,
+		stateChangeEvents:          make(chan statechange.Event),
+		ctx:                        context.TODO(),
+		engine: &DockerTaskEngine{
+			dataClient:          data.NewNoopClient(),
+			hostResourceManager: &hostResourceManager,
+		},
+	}
+	// Discard all the statechange events
+	defer discardEvents(mTask.stateChangeEvents)()
+
+	mTask.SetKnownStatus(apitaskstatus.TaskRunning)
+	mTask.SetSentStatus(apitaskstatus.TaskRunning)
+	container := mTask.Containers[0]
+
+	containerChange := dockerContainerChange{
+		container: container,
+		event: dockerapi.DockerContainerChangeEvent{
+			Status: apicontainerstatus.ContainerStopped,
+		},
+	}
+
+	mTask.handleContainerChange(containerChange)
+	// TODO validate what happens for a stopped container WITHOUT a restart policy
+}
+
+func TestHandleContainerChangeStopped_WithRestartPolicy(t *testing.T) {
+	eventStreamName := "TestHandleContainerChangeStoppedWithRestartPolicy"
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	containerChangeEventStream := eventstream.NewEventStream(eventStreamName, ctx)
+	containerChangeEventStream.StartListening()
+
+	hostResourceManager := NewHostResourceManager(getTestHostResources())
+	mTask := &managedTask{
+		Task:                       testdata.LoadTask("sleep5RestartPolicy"),
+		containerChangeEventStream: containerChangeEventStream,
+		stateChangeEvents:          make(chan statechange.Event),
+		ctx:                        context.TODO(),
+		engine: &DockerTaskEngine{
+			dataClient:          data.NewNoopClient(),
+			hostResourceManager: &hostResourceManager,
+		},
+	}
+	// Discard all the statechange events
+	defer discardEvents(mTask.stateChangeEvents)()
+
+	mTask.SetKnownStatus(apitaskstatus.TaskRunning)
+	mTask.SetSentStatus(apitaskstatus.TaskRunning)
+	container := mTask.Containers[0]
+	container.RestartTracker = restart.NewRestartTracker(*container.RestartPolicy)
+
+	exitCode := int(100)
+	containerChange := dockerContainerChange{
+		container: container,
+		event: dockerapi.DockerContainerChangeEvent{
+			Status: apicontainerstatus.ContainerStopped,
+			DockerContainerMetadata: dockerapi.DockerContainerMetadata{
+				ExitCode: &exitCode,
+			},
+	}
+
+	restartCount := container.RestartTracker.GetRestartCount()
+	assert.Equal(t, 0, restartCount, "Before stop event, restart count should be 0")
+	mTask.handleContainerChange(containerChange)
+	// wait for restart count to increment
+	for i := 0; i < 20; i++ {
+		restartCount = container.RestartTracker.GetRestartCount()
+		if restartCount == 1 {
+			break
+		}
+		time.Sleep(time.Millisecond * 500)
+	}
+	assert.Equal(t, 1, restartCount, "After stop event, container should have been restarted")
+}
+
+func TestHandleContainerChangeStopped_WithRestartPolicy_DesiredStopped(t *testing.T) {
+	eventStreamName := "TestHandleContainerChangeStoppedWithRestartPolicyDesiredStopped"
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	containerChangeEventStream := eventstream.NewEventStream(eventStreamName, ctx)
+	containerChangeEventStream.StartListening()
+
+	hostResourceManager := NewHostResourceManager(getTestHostResources())
+	mTask := &managedTask{
+		Task:                       testdata.LoadTask("sleep5RestartPolicy"),
+		containerChangeEventStream: containerChangeEventStream,
+		stateChangeEvents:          make(chan statechange.Event),
+		ctx:                        context.TODO(),
+		engine: &DockerTaskEngine{
+			dataClient:          data.NewNoopClient(),
+			hostResourceManager: &hostResourceManager,
+		},
+	}
+	// Discard all the statechange events
+	defer discardEvents(mTask.stateChangeEvents)()
+
+	mTask.SetKnownStatus(apitaskstatus.TaskRunning)
+	mTask.SetSentStatus(apitaskstatus.TaskRunning)
+	container := mTask.Containers[0]
+	container.RestartTracker = restart.NewRestartTracker(*container.RestartPolicy)
+
+	// TODO set desired stopped on task/containers
+
+	containerChange := dockerContainerChange{
+		container: container,
+		event: dockerapi.DockerContainerChangeEvent{
+			Status: apicontainerstatus.ContainerStopped,
+		},
+	}
+
+	mTask.handleContainerChange(containerChange)
+	// TODO validate what happens for a container that has stopped with a restart policy
+}
+
 func TestWaitForResourceTransition(t *testing.T) {
 	task := &managedTask{
 		Task: &apitask.Task{

agent/engine/testdata/test_tasks/sleep5RestartPolicy.json

@@ -0,0 +1,41 @@
+{
+    "Arn":"arn:aws:ecs:us-west-2:123456789012:task/12345678-90ab-cdef-1234-56780abcdef1",
+    "Family":"sleep5RestartPolicy",
+    "Version":"2",
+    "Containers":
+    [
+      {
+        "Name":"sleep5",
+        "Image":"busybox",
+        "Command":["sleep","5"],
+        "Cpu":10,
+        "Memory":10,
+        "Links":null,
+        "volumesFrom":[],
+        "mountPoints":[],
+        "portMappings":[],
+        "Essential":true,
+        "EntryPoint":null,
+        "environment":{},
+        "overrides":{"command":null},
+        "desiredStatus":"NONE",
+        "KnownStatus":"NONE",
+        "RunDependencies":null,
+        "IsInternal":false,
+        "AppliedStatus":"NONE",
+        "ApplyingError":null,
+        "SentStatus":"NONE",
+        "KnownExitCode":null,
+        "KnownPortBindingsUnsafe":null,
+        "StatusLock":{},
+        "RestartPolicy": {
+            "Enabled":true
+        }
+      }
+    ],
+    "volumes":[],
+    "DesiredStatus":"RUNNING",
+    "KnownStatus":"NONE",
+    "KnownTime":"0001-01-01T00:00:00Z",
+    "SentStatus":"NONE"
+  }