Consume ECS client from ecs-agent module in agent module

aws · Nov 20, 2023 · 5f35c93 · 5f35c93
1 parent a540013
commit 5f35c93
Show file tree

Hide file tree

Showing 87 changed files with 1,650 additions and 4,044 deletions.
diff --git a/agent/acs/session/payload_responder.go b/agent/acs/session/payload_responder.go
@@ -23,6 +23,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/eventhandler"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
 	apiresource "github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
@@ -40,7 +41,7 @@ type skipAddTaskComparatorFunc func(apitaskstatus.TaskStatus) bool
 // payloadMessageHandler implements PayloadMessageHandler interface defined in ecs-agent module.
 type payloadMessageHandler struct {
 	taskEngine                  engine.TaskEngine
-	ecsClient                   api.ECSClient
+	ecsClient                   ecs.ECSClient
 	dataClient                  data.Client
 	taskHandler                 *eventhandler.TaskHandler
 	credentialsManager          credentials.Manager
@@ -49,7 +50,7 @@ type payloadMessageHandler struct {
 
 // NewPayloadMessageHandler creates a new payloadMessageHandler.
 func NewPayloadMessageHandler(taskEngine engine.TaskEngine,
-	ecsClient api.ECSClient,
+	ecsClient ecs.ECSClient,
 	dataClient data.Client,
 	taskHandler *eventhandler.TaskHandler,
 	credentialsManager credentials.Manager,

diff --git a/agent/acs/session/payload_responder_test.go b/agent/acs/session/payload_responder_test.go
@@ -22,8 +22,6 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/aws/amazon-ecs-agent/agent/api"
-	mock_api "github.com/aws/amazon-ecs-agent/agent/api/mocks"
 	apitask "github.com/aws/amazon-ecs-agent/agent/api/task"
 	"github.com/aws/amazon-ecs-agent/agent/data"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
@@ -35,6 +33,8 @@ import (
 	acssession "github.com/aws/amazon-ecs-agent/ecs-agent/acs/session"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/session/testconst"
 	apiresource "github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
+	mock_ecs "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
@@ -76,7 +76,7 @@ type testHelper struct {
 func setup(t *testing.T, acsResponseSender wsclient.RespondFunc) *testHelper {
 	ctrl := gomock.NewController(t)
 	taskEngine := mock_engine.NewMockTaskEngine(ctrl)
-	ecsClient := mock_api.NewMockECSClient(ctrl)
+	ecsClient := mock_ecs.NewMockECSClient(ctrl)
 	dataClient := data.NewNoopClient()
 	credentialsManager := credentials.NewManager()
 	ctx := context.Background()
@@ -1057,7 +1057,7 @@ func TestHandlePayloadMessageAddedFirelensData(t *testing.T) {
 
 func TestHandleInvalidTask(t *testing.T) {
 	tester := setup(t, nil)
-	mockECSACSClient := mock_api.NewMockECSClient(tester.ctrl)
+	mockECSACSClient := mock_ecs.NewMockECSClient(tester.ctrl)
 	taskHandler := eventhandler.NewTaskHandler(tester.ctx, data.NewNoopClient(), dockerstate.NewTaskEngineState(),
 		mockECSACSClient)
 	tester.payloadMessageHandler.ecsClient = mockECSACSClient
@@ -1070,8 +1070,8 @@ func TestHandleInvalidTask(t *testing.T) {
 	wait := &sync.WaitGroup{}
 	wait.Add(1)
 
-	mockECSACSClient.EXPECT().SubmitTaskStateChange(gomock.Any()).Do(func(change api.TaskStateChange) {
-		assert.NotNil(t, change.Task)
+	mockECSACSClient.EXPECT().SubmitTaskStateChange(gomock.Any()).Do(func(change ecs.TaskStateChange) {
+		assert.False(t, change.MetadataGetter.GetTaskIsNil())
 		wait.Done()
 	})