more code review fixups

aws · Sep 12, 2019 · 272aafb · 272aafb
1 parent 0dc4b86
commit 272aafb
Show file tree

Hide file tree

Showing 2 changed files with 107 additions and 5 deletions.
diff --git a/agent/app/agent.go b/agent/app/agent.go
@@ -15,6 +15,7 @@ package app
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"time"
@@ -630,9 +631,28 @@ func (agent *ecsAgent) startSpotInstanceDrainingPoller(client api.ECSClient) {
 // set AND the container instance state is successfully updated to DRAINING.
 func (agent *ecsAgent) spotInstanceDrainingPoller(client api.ECSClient) bool {
 	// this endpoint 404s unless a interruption has been set, so expect failure in most cases.
-	termtime, err := agent.ec2MetadataClient.SpotInstanceAction()
-	if err == nil && len(termtime) > 0 {
-		seelog.Infof("Received a spot interruption (%s), setting state to DRAINING", termtime)
+	tmp, err := agent.ec2MetadataClient.SpotInstanceAction()
+	if err == nil {
+		type InstanceAction struct {
+			Time   string
+			Action string
+		}
+		ia := InstanceAction{}
+
+		err := json.Unmarshal([]byte(tmp), &ia)
+		if err != nil {
+			seelog.Errorf("Invalid response from /spot/instance-action endpoint: %s Error: %s", tmp, err)
+			return false
+		}
+
+		switch ia.Action {
+		case "hibernate", "terminate", "stop":
+		default:
+			seelog.Errorf("Invalid response from /spot/instance-action endpoint: %s, Error: unrecognized action (%s)", tmp, ia.Action)
+			return false
+		}
+
+		seelog.Infof("Received a spot interruption (%s) scheduled for %s, setting state to DRAINING", ia.Action, ia.Time)
 		err = client.UpdateContainerInstancesState(agent.containerInstanceARN, "DRAINING")
 		if err != nil {
 			seelog.Errorf("Error setting instance [ARN: %s] state to DRAINING: %s", agent.containerInstanceARN, err)

diff --git a/agent/app/agent_test.go b/agent/app/agent_test.go
@@ -1182,7 +1182,7 @@ func TestGetHostPublicIPv4AddressFromEC2MetadataFailWithError(t *testing.T) {
 	assert.Empty(t, agent.getHostPublicIPv4AddressFromEC2Metadata())
 }
 
-func TestSpotInstanceActionCheck_Yes(t *testing.T) {
+func TestSpotInstanceActionCheck_Terminate(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 
@@ -1202,7 +1202,47 @@ func TestSpotInstanceActionCheck_Yes(t *testing.T) {
 	assert.True(t, agent.spotInstanceDrainingPoller(ecsClient))
 }
 
-func TestSpotInstanceActionCheck_EmptyTimestamp(t *testing.T) {
+func TestSpotInstanceActionCheck_Stop(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	ec2MetadataClient := mock_ec2.NewMockEC2MetadataClient(ctrl)
+	ec2Client := mock_ec2.NewMockClient(ctrl)
+	ecsClient := mock_api.NewMockECSClient(ctrl)
+
+	myARN := "myARN"
+	agent := &ecsAgent{
+		ec2MetadataClient:    ec2MetadataClient,
+		ec2Client:            ec2Client,
+		containerInstanceARN: myARN,
+	}
+	ec2MetadataClient.EXPECT().SpotInstanceAction().Return("{\"action\": \"stop\", \"time\": \"2017-09-18T08:22:00Z\"}", nil)
+	ecsClient.EXPECT().UpdateContainerInstancesState(myARN, "DRAINING").Return(nil)
+
+	assert.True(t, agent.spotInstanceDrainingPoller(ecsClient))
+}
+
+func TestSpotInstanceActionCheck_Hibernate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	ec2MetadataClient := mock_ec2.NewMockEC2MetadataClient(ctrl)
+	ec2Client := mock_ec2.NewMockClient(ctrl)
+	ecsClient := mock_api.NewMockECSClient(ctrl)
+
+	myARN := "myARN"
+	agent := &ecsAgent{
+		ec2MetadataClient:    ec2MetadataClient,
+		ec2Client:            ec2Client,
+		containerInstanceARN: myARN,
+	}
+	ec2MetadataClient.EXPECT().SpotInstanceAction().Return("{\"action\": \"hibernate\", \"time\": \"2017-09-18T08:22:00Z\"}", nil)
+	ecsClient.EXPECT().UpdateContainerInstancesState(myARN, "DRAINING").Return(nil)
+
+	assert.True(t, agent.spotInstanceDrainingPoller(ecsClient))
+}
+
+func TestSpotInstanceActionCheck_EmptyJSON(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 
@@ -1223,6 +1263,48 @@ func TestSpotInstanceActionCheck_EmptyTimestamp(t *testing.T) {
 	assert.False(t, agent.spotInstanceDrainingPoller(ecsClient))
 }
 
+func TestSpotInstanceActionCheck_InvalidJSON(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	ec2MetadataClient := mock_ec2.NewMockEC2MetadataClient(ctrl)
+	ec2Client := mock_ec2.NewMockClient(ctrl)
+	ecsClient := mock_api.NewMockECSClient(ctrl)
+
+	myARN := "myARN"
+	agent := &ecsAgent{
+		ec2MetadataClient:    ec2MetadataClient,
+		ec2Client:            ec2Client,
+		containerInstanceARN: myARN,
+	}
+	ec2MetadataClient.EXPECT().SpotInstanceAction().Return("{\"action\": \"terminate\" \"time\": \"2017-09-18T08:22:00Z\"}", nil)
+	// Container state should NOT be updated because the termination time field is empty.
+	ecsClient.EXPECT().UpdateContainerInstancesState(gomock.Any(), gomock.Any()).Times(0)
+
+	assert.False(t, agent.spotInstanceDrainingPoller(ecsClient))
+}
+
+func TestSpotInstanceActionCheck_UnknownInstanceAction(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	ec2MetadataClient := mock_ec2.NewMockEC2MetadataClient(ctrl)
+	ec2Client := mock_ec2.NewMockClient(ctrl)
+	ecsClient := mock_api.NewMockECSClient(ctrl)
+
+	myARN := "myARN"
+	agent := &ecsAgent{
+		ec2MetadataClient:    ec2MetadataClient,
+		ec2Client:            ec2Client,
+		containerInstanceARN: myARN,
+	}
+	ec2MetadataClient.EXPECT().SpotInstanceAction().Return("{\"action\": \"flip!\", \"time\": \"2017-09-18T08:22:00Z\"}", nil)
+	// Container state should NOT be updated because the termination time field is empty.
+	ecsClient.EXPECT().UpdateContainerInstancesState(gomock.Any(), gomock.Any()).Times(0)
+
+	assert.False(t, agent.spotInstanceDrainingPoller(ecsClient))
+}
+
 func TestSpotInstanceActionCheck_No(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()