unit test

agent/api/container/container_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/utils"
 	dockercontainer "github.com/docker/docker/api/types/container"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type configPair struct {
@@ -1334,6 +1335,42 @@ func TestGetCredentialSpec(t *testing.T) {
 	}
 }
 
+func TestRestartPolicy(t *testing.T) {
+	testCases := []struct {
+		name            string
+		container       *Container
+		restartCount    int
+		expectedEnabled bool
+	}{
+		{
+			name: "base case",
+			container: &Container{
+				RestartPolicy: RestartPolicy{},
+			},
+			restartCount:    0,
+			expectedEnabled: false,
+		},
+		{
+			name: "do 5 restarts",
+			container: &Container{
+				RestartPolicy: RestartPolicy{
+					Enabled: true,
+				},
+			},
+			restartCount:    5,
+			expectedEnabled: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		require.Equal(t, tc.expectedEnabled, tc.container.RestartPolicyEnabled())
+		for i := 0; i < tc.restartCount; i++ {
+			tc.container.RecordRestart()
+		}
+		require.Equal(t, tc.restartCount, tc.container.GetRestartCount())
+	}
+}
+
 func getContainer(hostConfig string, credentialSpecs []string) *Container {
 	c := &Container{
 		Name: "c",