acs: return an empty task when task is not recognized

aws · Aug 9, 2018 · 17809f3 · 17809f3
1 parent dfbe8cc
commit 17809f3
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 4 deletions.
diff --git a/agent/acs/handler/acs_handler.go b/agent/acs/handler/acs_handler.go
@@ -16,14 +16,13 @@
 package handler
 
 import (
+	"context"
 	"io"
 	"net/url"
 	"strconv"
 	"strings"
 	"time"
 
-	"context"
-
 	acsclient "github.com/aws/amazon-ecs-agent/agent/acs/client"
 	"github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs"
 	"github.com/aws/amazon-ecs-agent/agent/acs/update_handler"

diff --git a/agent/acs/handler/payload_handler.go b/agent/acs/handler/payload_handler.go
@@ -337,6 +337,10 @@ func (payloadHandler *payloadRequestHandler) handleUnrecognizedTask(task *ecsacs
 		TaskARN: *task.Arn,
 		Status:  apitaskstatus.TaskStopped,
 		Reason:  UnrecognizedTaskError{err}.Error(),
+		// The real task cannot be extracted from payload message, so we send an empty task.
+		// This is necessary because the task handler will not send an event whose
+		// Task is nil.
+		Task: &apitask.Task{},
 	}
 
 	payloadHandler.taskHandler.AddStateChangeEvent(taskEvent, payloadHandler.ecsClient)

diff --git a/agent/acs/handler/payload_handler_test.go b/agent/acs/handler/payload_handler_test.go
@@ -15,12 +15,13 @@
 package handler
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"reflect"
+	"sync"
 	"testing"
 
-	"context"
-
 	"github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs"
 	"github.com/aws/amazon-ecs-agent/agent/api"
 	"github.com/aws/amazon-ecs-agent/agent/api/mocks"
@@ -754,3 +755,31 @@ func TestPayloadHandlerAddedASMAuthData(t *testing.T) {
 	assert.Equal(t, aws.StringValue(expected.AsmAuthData.Region), actual.ASMAuthData.Region)
 	assert.Equal(t, aws.StringValue(expected.AsmAuthData.CredentialsParameter), actual.ASMAuthData.CredentialsParameter)
 }
+
+func TestHandleUnrecognizedTask(t *testing.T) {
+	tester := setup(t)
+	defer tester.ctrl.Finish()
+
+	arn := "arn"
+	ecsacsTask := &ecsacs.Task{Arn: &arn}
+	payloadMessage := &ecsacs.PayloadMessage{
+		Tasks: []*ecsacs.Task{ecsacsTask},
+		MessageId: aws.String(payloadMessageId),
+	}
+
+	mockECSACSClient := mock_api.NewMockECSClient(tester.ctrl)
+	taskHandler := eventhandler.NewTaskHandler(tester.ctx, tester.payloadHandler.saver, nil, mockECSACSClient)
+	tester.payloadHandler.taskHandler = taskHandler
+
+	wait := &sync.WaitGroup{}
+	wait.Add(1)
+
+	mockECSACSClient.EXPECT().SubmitTaskStateChange(gomock.Any()).Do(func(change api.TaskStateChange) {
+		assert.NotNil(t, change.Task)
+		wait.Done()
+		// Don't return, verify timeout happens.
+	})
+
+	tester.payloadHandler.handleUnrecognizedTask(ecsacsTask, errors.New("test error"), payloadMessage)
+	wait.Wait()
+}