From 272aafbc7cfadcd90302f17a82d9ff6f3cacb2dc Mon Sep 17 00:00:00 2001 From: Cam Date: Thu, 12 Sep 2019 16:19:46 -0700 Subject: [PATCH] more code review fixups --- agent/app/agent.go | 26 +++++++++++-- agent/app/agent_test.go | 86 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 107 insertions(+), 5 deletions(-) diff --git a/agent/app/agent.go b/agent/app/agent.go index f62084af5ab..999568e7a1e 100644 --- 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 index e17d6d89c0f..82f91d48f44 100644 --- 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()