test

loilo-inc · Jul 3, 2024 · 14d22d7 · 14d22d7
1 parent 305c2f2
commit 14d22d7
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/task/alb_task.go b/task/alb_task.go
@@ -144,7 +144,7 @@ func (c *albTask) getEc2TargetNetwork(ctx context.Context) (*string, *string, er
 	} else {
 		ec2Instance = o.Reservations[0].Instances[0]
 	}
-	return ec2Instance.PrivateIpAddress, ec2Instance.SubnetId, nil
+	return ec2Instance.InstanceId, ec2Instance.SubnetId, nil
 }
 
 func (c *albTask) getTargetPort() (*int32, error) {

diff --git a/task/alb_task_test.go b/task/alb_task_test.go
@@ -228,7 +228,13 @@ func TestAlbTask_RegisterToTargetGroup(t *testing.T) {
 			ec2Mock.EXPECT().DescribeSubnets(gomock.Any(), gomock.Any()).Return(&ec2.DescribeSubnetsOutput{
 				Subnets: subnets,
 			}, nil)
-			albMock.EXPECT().RegisterTargets(gomock.Any(), gomock.Any()).Return(nil, nil)
+			albMock.EXPECT().RegisterTargets(gomock.Any(), &elbv2.RegisterTargetsInput{
+				TargetGroupArn: atask.Lb.TargetGroupArn,
+				Targets: []elbv2types.TargetDescription{{
+					Id:               aws.String("127.0.0.1"),
+					Port:             aws.Int32(80),
+					AvailabilityZone: subnets[0].AvailabilityZone},
+				}}).Return(nil, nil)
 			atask.TaskArn = aws.String("arn://task")
 			err := atask.RegisterToTargetGroup(context.TODO())
 			assert.NoError(t, err)
@@ -278,10 +284,13 @@ func TestAlbTask_RegisterToTargetGroup(t *testing.T) {
 	t.Run("EC2", func(t *testing.T) {
 		containerInstances := []ecstypes.ContainerInstance{{
 			ContainerInstanceArn: aws.String("arn://container"),
+			Ec2InstanceId:        aws.String("i-123456"),
 		}}
 		reservations := []ec2types.Reservation{{
 			Instances: []ec2types.Instance{{
-				InstanceId: aws.String("i-123456"),
+				InstanceId:       aws.String("i-123456"),
+				SubnetId:         aws.String("subnet-123456"),
+				PrivateIpAddress: aws.String("127.0.0.1"),
 			}},
 		}}
 		subnets := []ec2types.Subnet{{
@@ -290,6 +299,7 @@ func TestAlbTask_RegisterToTargetGroup(t *testing.T) {
 		setup := func(t *testing.T) (*mock_awsiface.MockEc2Client, *mock_awsiface.MockAlbClient, *mock_awsiface.MockEcsClient, *task.AlbTaskExport) {
 			ctrl := gomock.NewController(t)
 			envars := test.DefaultEnvars()
+			envars.CanaryInstanceArn = "arn://container"
 			mocker := test.NewMockContext()
 			td, _ := mocker.Ecs.RegisterTaskDefinition(context.TODO(), envars.TaskDefinitionInput)
 			ec2Mock := mock_awsiface.NewMockEc2Client(ctrl)
@@ -318,7 +328,13 @@ func TestAlbTask_RegisterToTargetGroup(t *testing.T) {
 			ec2Mock.EXPECT().DescribeSubnets(gomock.Any(), gomock.Any()).Return(&ec2.DescribeSubnetsOutput{
 				Subnets: subnets,
 			}, nil)
-			albMock.EXPECT().RegisterTargets(gomock.Any(), gomock.Any()).Return(nil, nil)
+			albMock.EXPECT().RegisterTargets(gomock.Any(), &elbv2.RegisterTargetsInput{
+				TargetGroupArn: atask.Lb.TargetGroupArn,
+				Targets: []elbv2types.TargetDescription{{
+					Id:               containerInstances[0].Ec2InstanceId,
+					Port:             aws.Int32(80),
+					AvailabilityZone: subnets[0].AvailabilityZone},
+				}}).Return(nil, nil)
 			err := atask.RegisterToTargetGroup(context.TODO())
 			assert.NoError(t, err)
 		})

diff --git a/test/ecs.go b/test/ecs.go
@@ -168,6 +168,8 @@ func (ctx *EcsServer) StartTask(_ context.Context, input *ecs.StartTaskInput, _
 	}
 	taskArn := fmt.Sprintf("arn:aws:ecs:us-west-2:012345678910:task/%s", uuid.New().String())
 	attachment := types.Attachment{
+		Status: aws.String("ATTACHED"),
+		Type:   aws.String("ElasticNetworkInterface"),
 		Details: []types.KeyValuePair{
 			{
 				Name:  aws.String("privateIPv4Address"),