Update alb_task_test.go

loilo-inc · Jul 3, 2024 · 730399f · 730399f
1 parent 3172b6e
commit 730399f
Showing 1 changed file with 47 additions and 25 deletions.
diff --git a/task/alb_task_test.go b/task/alb_task_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
+	"github.com/loilo-inc/canarycage/env"
 	"github.com/loilo-inc/canarycage/key"
 	"github.com/loilo-inc/canarycage/task"
 	"github.com/loilo-inc/canarycage/test"
@@ -13,31 +14,52 @@ import (
 )
 
 func TestAlbTask(t *testing.T) {
-	mocker := test.NewMockContext()
-	env := test.DefaultEnvars()
-	ctx := context.TODO()
-	td, _ := mocker.Ecs.RegisterTaskDefinition(ctx, env.TaskDefinitionInput)
-	env.ServiceDefinitionInput.TaskDefinition = td.TaskDefinition.TaskDefinitionArn
-	ecsSvc, _ := mocker.Ecs.CreateService(ctx, env.ServiceDefinitionInput)
-	d := di.NewDomain(func(b *di.B) {
-		b.Set(key.Env, env)
-		b.Set(key.EcsCli, mocker.Ecs)
-		b.Set(key.Ec2Cli, mocker.Ec2)
-		b.Set(key.AlbCli, mocker.Alb)
-		b.Set(key.Time, test.NewFakeTime())
+	setup := func(env *env.Envars) (task.Task, *test.MockContext) {
+		mocker := test.NewMockContext()
+		ctx := context.TODO()
+		td, _ := mocker.Ecs.RegisterTaskDefinition(ctx, env.TaskDefinitionInput)
+		env.ServiceDefinitionInput.TaskDefinition = td.TaskDefinition.TaskDefinitionArn
+		ecsSvc, _ := mocker.Ecs.CreateService(ctx, env.ServiceDefinitionInput)
+		d := di.NewDomain(func(b *di.B) {
+			b.Set(key.Env, env)
+			b.Set(key.EcsCli, mocker.Ecs)
+			b.Set(key.Ec2Cli, mocker.Ec2)
+			b.Set(key.AlbCli, mocker.Alb)
+			b.Set(key.Time, test.NewFakeTime())
+		})
+		stask := task.NewAlbTask(d, &task.Input{
+			TaskDefinition:       td.TaskDefinition,
+			NetworkConfiguration: ecsSvc.Service.NetworkConfiguration,
+		}, &ecsSvc.Service.LoadBalancers[0])
+		mocker.Alb.RegisterTargets(ctx, &elbv2.RegisterTargetsInput{
+			TargetGroupArn: ecsSvc.Service.LoadBalancers[0].TargetGroupArn,
+		})
+		return stask, mocker
+	}
+	t.Run("fargate", func(t *testing.T) {
+		env := test.DefaultEnvars()
+		stask, mocker := setup(env)
+		ctx := context.TODO()
+		err := stask.Start(ctx)
+		assert.NoError(t, err)
+		err = stask.Wait(ctx)
+		assert.NoError(t, err)
+		err = stask.Stop(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, 1, mocker.RunningTaskSize())
 	})
-	stask := task.NewAlbTask(d, &task.Input{
-		TaskDefinition:       td.TaskDefinition,
-		NetworkConfiguration: ecsSvc.Service.NetworkConfiguration,
-	}, &ecsSvc.Service.LoadBalancers[0])
-	mocker.Alb.RegisterTargets(ctx, &elbv2.RegisterTargetsInput{
-		TargetGroupArn: ecsSvc.Service.LoadBalancers[0].TargetGroupArn,
+	t.Run("ec2", func(t *testing.T) {
+		env := test.DefaultEnvars()
+		env.CanaryInstanceArn = "arn://ec2"
+		stask, mocker := setup(env)
+		ctx := context.TODO()
+		err := stask.Start(ctx)
+		assert.NoError(t, err)
+		err = stask.Wait(ctx)
+		assert.NoError(t, err)
+		err = stask.Stop(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, 1, mocker.RunningTaskSize())
 	})
-	err := stask.Start(ctx)
-	assert.NoError(t, err)
-	err = stask.Wait(ctx)
-	assert.NoError(t, err)
-	err = stask.Stop(ctx)
-	assert.NoError(t, err)
-	assert.Equal(t, 1, mocker.RunningTaskSize())
+
 }