Implement TextMarshaler for ContainerStatus

agent/api/container/transitiondependency_test.go

@@ -20,42 +20,90 @@ import (
 	"encoding/json"
 	"testing"
 
+	"github.com/aws/amazon-ecs-agent/agent/taskresource/status"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-func TestUnmarshalOldTransitionDependencySet(t *testing.T) {
-	bytes := []byte(`{
-	  "ContainerDependencies": [
-	    {
-	      "ContainerName": "container",
-	      "SatisfiedStatus": "RUNNING",
-	      "DependentStatus": "RUNNING"
-	    }
-	  ]
-	}`)
-	unmarshalledTdMap := TransitionDependenciesMap{}
-	err := json.Unmarshal(bytes, &unmarshalledTdMap)
-	assert.NoError(t, err)
-	assert.Len(t, unmarshalledTdMap, 1)
-	assert.NotNil(t, unmarshalledTdMap[apicontainerstatus.ContainerRunning])
-	dep := unmarshalledTdMap[apicontainerstatus.ContainerRunning].ContainerDependencies
-	assert.Len(t, dep, 1)
-	assert.Equal(t, "container", dep[0].ContainerName)
-	assert.Equal(t, apicontainerstatus.ContainerRunning, dep[0].SatisfiedStatus)
-	assert.Equal(t, apicontainerstatus.ContainerStatusNone, dep[0].DependentStatus)
+// Tests that TransitionDependencySet marshaled in different formats (as changes are introduced)
+// can all be unmarshaled.
+func TestUnmarshalTransitionDependencySet(t *testing.T) {
+	tcs := []struct {
+		name      string
+		marshaled []byte
+	}{
+		{
+			name:      "with dependent status - this is the oldest format",
+			marshaled: []byte(`{"ContainerDependencies": [{"ContainerName": "container", "SatisfiedStatus": "RUNNING", "DependentStatus": "PULLED"}]}`),
+		},
+		{
+			name:      "as a map with integer keys for statuses - this is an older format",
+			marshaled: []byte(`{"1":{"ContainerDependencies":[{"ContainerName":"container","SatisfiedStatus":"RUNNING"}]}}`),
+		},
+		{
+			name:      "as a map with string keys for statuses - this is the latest format",
+			marshaled: []byte(`{"PULLED":{"ContainerDependencies":[{"ContainerName":"container","SatisfiedStatus":"RUNNING"}]}}`),
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			unmarshalledTdMap := TransitionDependenciesMap{}
+			err := json.Unmarshal(tc.marshaled, &unmarshalledTdMap)
+			require.NoError(t, err)
+			require.Len(t, unmarshalledTdMap, 1)
+			require.NotNil(t, unmarshalledTdMap[apicontainerstatus.ContainerPulled])
+			dep := unmarshalledTdMap[apicontainerstatus.ContainerPulled].ContainerDependencies
+			require.Len(t, dep, 1)
+			assert.Equal(t, "container", dep[0].ContainerName)
+			assert.Equal(t, apicontainerstatus.ContainerRunning, dep[0].SatisfiedStatus)
+			assert.Equal(t, apicontainerstatus.ContainerStatusNone, dep[0].DependentStatus)
+		})
+	}
 }
 
-func TestUnmarshalNewTransitionDependencySet(t *testing.T) {
-	bytes := []byte(`{"1":{"ContainerDependencies":[{"ContainerName":"container","SatisfiedStatus":"RUNNING"}]}}`)
-	unmarshalledTdMap := TransitionDependenciesMap{}
-	err := json.Unmarshal(bytes, &unmarshalledTdMap)
-	assert.NoError(t, err)
-	assert.Len(t, unmarshalledTdMap, 1)
-	assert.NotNil(t, unmarshalledTdMap[apicontainerstatus.ContainerPulled])
-	dep := unmarshalledTdMap[apicontainerstatus.ContainerPulled].ContainerDependencies
-	assert.Len(t, dep, 1)
-	assert.Equal(t, "container", dep[0].ContainerName)
-	assert.Equal(t, apicontainerstatus.ContainerRunning, dep[0].SatisfiedStatus)
-	assert.Equal(t, apicontainerstatus.ContainerStatusNone, dep[0].DependentStatus)
+// Tests that marshaled TransitionDependenciesMap can be unmarshaled.
+func TestMarshalUnmarshalTransitionDependencySet(t *testing.T) {
+	var depMap TransitionDependenciesMap = map[apicontainerstatus.ContainerStatus]TransitionDependencySet{
+		apicontainerstatus.ContainerRunning: {
+			ContainerDependencies: []ContainerDependency{
+				{
+					ContainerName:   "db",
+					SatisfiedStatus: apicontainerstatus.ContainerRunning,
+				},
+			},
+			ResourceDependencies: []ResourceDependency{
+				{
+					Name:           "config",
+					RequiredStatus: status.ResourceCreated,
+				},
+			},
+		},
+	}
+	marshaled, err := json.Marshal(depMap)
+	require.NoError(t, err)
+	var unmarshaled TransitionDependenciesMap
+	err = json.Unmarshal(marshaled, &unmarshaled)
+	require.NoError(t, err)
+	assert.Equal(t, depMap, unmarshaled)
+}
+
+// Tests that marshaling of TransitionDependenciesMap works as expected.
+func TestMarshalTransitionDependencySet(t *testing.T) {
+	var depMap TransitionDependenciesMap = map[apicontainerstatus.ContainerStatus]TransitionDependencySet{
+		apicontainerstatus.ContainerPulled: {
+			ContainerDependencies: []ContainerDependency{
+				{
+					ContainerName:   "pre",
+					SatisfiedStatus: apicontainerstatus.ContainerStopped,
+				},
+			},
+		},
+	}
+	marshaled, err := json.Marshal(depMap)
+	require.NoError(t, err)
+	assert.Equal(
+		t,
+		`{"PULLED":{"ContainerDependencies":[{"ContainerName":"pre","SatisfiedStatus":"STOPPED"}],"ResourceDependencies":null}}`,
+		string(marshaled))
 }

ecs-agent/api/container/status/containerstatus.go

@@ -15,6 +15,7 @@ package status
 
 import (
 	"errors"
+	"strconv"
 	"strings"
 )
 
@@ -138,6 +139,28 @@ func (cs *ContainerStatus) UnmarshalJSON(b []byte) error {
 		return nil
 	}
 
+	// Before MarshalText method was implemented for ContainerStatus, ContainerStatus was
+	// marshaled as an integer in cases that depend on MarshalText such as keys in a JSON object.
+	// To make unmarshaling backwards-compatible, check if the marshaled text is an integer
+	// and map it to container status.
+	if intStatus, err := strconv.Atoi(strStatus); err == nil {
+		// This map is only for making text unmarshaling compatible with old state files.
+		// Updates are NOT needed to this map as newer container states are introduced.
+		intContainerStatusMap := map[int]ContainerStatus{
+			0: ContainerStatusNone,
+			1: ContainerPulled,
+			2: ContainerCreated,
+			3: ContainerRunning,
+			4: ContainerResourcesProvisioned,
+			5: ContainerStopped,
+			6: ContainerZombie,
+		}
+		if stat, ok := intContainerStatusMap[intStatus]; ok {
+			*cs = stat
+			return nil
+		}
+	}
+
 	stat, ok := containerStatusMap[strStatus]
 	if !ok {
 		*cs = ContainerStatusNone
@@ -155,6 +178,26 @@ func (cs *ContainerStatus) MarshalJSON() ([]byte, error) {
 	return []byte(`"` + cs.String() + `"`), nil
 }
 
+// Marshals a container status to its text form.
+// In some cases such as a map with ContainerStatus as keys, MarshalText method is used to
+// marshal container statuses by functions in the encoding package. Without this method
+// container statuses will be marshaled as integers which is undesirable.
+func (cs ContainerStatus) MarshalText() ([]byte, error) {
+	return []byte(cs.String()), nil
+}
+
+// Unmarshals a container status from its text form.
+func (cs *ContainerStatus) UnmarshalText(b []byte) error {
+	strStatus := string(b)
+	stat, ok := containerStatusMap[strStatus]
+	if !ok {
+		*cs = ContainerStatusNone
+		return errors.New("container status text unmarshal: unrecognized status: " + strStatus)
+	}
+	*cs = stat
+	return nil
+}
+
 // UnmarshalJSON overrides the logic for parsing the JSON-encoded container health data
 func (healthStatus *ContainerHealthStatus) UnmarshalJSON(b []byte) error {
 	*healthStatus = ContainerHealthUnknown

ecs-agent/api/container/status/containerstatus_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestShouldReportToBackend(t *testing.T) {
@@ -171,3 +172,114 @@ func TestUnmarshalContainerHealthStatus(t *testing.T) {
 		})
 	}
 }
+
+// Tests that all container statuses are marshaled to JSON with a quoted string.
+// Also tests that JSON marshaled container status can be unmarshaled.
+func TestContainerStatusMarshalUnmarshalJSON(t *testing.T) {
+	for strStatus, status := range containerStatusMap {
+		t.Run(fmt.Sprintf("marshal-unmarshal %v", strStatus), func(t *testing.T) {
+			marshaled, err := json.Marshal(status)
+			require.NoError(t, err)
+			require.Equal(t, fmt.Sprintf("%q", strStatus), string(marshaled))
+
+			var unmarshaled ContainerStatus
+			err = json.Unmarshal(marshaled, &unmarshaled)
+			require.NoError(t, err)
+			require.Equal(t, status, unmarshaled)
+		})
+	}
+}
+
+// Tests that a container status marshaled as text can be unmarshaled.
+func TestContainerStatusMarshalUnmarshalText(t *testing.T) {
+	for strStatus, status := range containerStatusMap {
+		t.Run(fmt.Sprintf("marshal-unmarshal %v", strStatus), func(t *testing.T) {
+			marshaled, err := status.MarshalText()
+			require.NoError(t, err)
+			require.Equal(t, fmt.Sprintf("%s", strStatus), string(marshaled))
+
+			var unmarshaled ContainerStatus
+			err = unmarshaled.UnmarshalText(marshaled)
+			require.NoError(t, err)
+			require.Equal(t, status, unmarshaled)
+		})
+	}
+}
+
+// Tests that MarshalText works as expected for container status pointers.
+func TestContainerStatusMarshalPointer(t *testing.T) {
+	status := ContainerPulled
+	ptr := &status
+	marshaled, err := ptr.MarshalText()
+	require.NoError(t, err)
+	assert.Equal(t, "PULLED", string(marshaled))
+}
+
+// Tests that unmarshaling an invalid text to container status fails.
+func TestContainerStatusTextUnmarshalError(t *testing.T) {
+	var status ContainerStatus
+	assert.EqualError(t, status.UnmarshalText([]byte("invalidStatus")),
+		"container status text unmarshal: unrecognized status: invalidStatus")
+}
+
+// Tests that string based statuses are used when a map with container status as keys is
+// marshaled to JSON.
+func TestContainerStatusKeyMarshal(t *testing.T) {
+	someMap := map[ContainerStatus]string{
+		ContainerStatusNone:           "",
+		ContainerPulled:               "",
+		ContainerCreated:              "",
+		ContainerRunning:              "",
+		ContainerResourcesProvisioned: "",
+		ContainerStopped:              "",
+	}
+	marshaled, err := json.Marshal(someMap)
+	require.NoError(t, err)
+
+	var unmarshaledMap map[string]string
+	err = json.Unmarshal(marshaled, &unmarshaledMap)
+	require.NoError(t, err)
+
+	assert.Equal(t, map[string]string{
+		`NONE`:                  "",
+		`PULLED`:                "",
+		`CREATED`:               "",
+		`RUNNING`:               "",
+		`RESOURCES_PROVISIONED`: "",
+		`STOPPED`:               "",
+	}, unmarshaledMap)
+
+	var unmarshaled map[ContainerStatus]string
+	err = json.Unmarshal(marshaled, &unmarshaled)
+	require.NoError(t, err)
+	assert.Equal(t, someMap, unmarshaled)
+}
+
+// Tests that JSON unmarshal of container status is backwards-compatible with legacy integer
+// based representations for JSON object keys.
+func TestContainerStatusJSONUnmarshalInt(t *testing.T) {
+	tcs := map[string]ContainerStatus{
+		`"0"`: ContainerStatusNone,
+		`"1"`: ContainerPulled,
+		`"2"`: ContainerCreated,
+		`"3"`: ContainerRunning,
+		`"4"`: ContainerResourcesProvisioned,
+		`"5"`: ContainerStopped,
+		`"6"`: ContainerZombie,
+	}
+	for intStatus, status := range tcs {
+		t.Run(fmt.Sprintf("%s - %s", intStatus, status.String()), func(t *testing.T) {
+			var unmarshaled ContainerStatus
+			err := json.Unmarshal([]byte(intStatus), &unmarshaled)
+			require.NoError(t, err)
+			assert.Equal(t, status, unmarshaled)
+		})
+	}
+}
+
+func TestTemporary(t *testing.T) {
+	marshaled := `{"1": "ok"}`
+	unmarshaled := map[ContainerStatus]string{}
+	err := json.Unmarshal([]byte(marshaled), &unmarshaled)
+	require.NoError(t, err)
+}