Fix regexp for ECS_CLUSTER value fetching

.github/workflows/test.yml

@@ -0,0 +1,28 @@
+name: Test
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    branches:
+      - master
+jobs:
+  test:
+    strategy:
+      fail-fast: false
+      matrix:
+        go-version: [1.15.x, 1.16.x, 1.17.x]
+        platform: [ubuntu-latest]
+    runs-on: ${{ matrix.platform }}
+    steps:
+
+    - name: Install Go
+      uses: actions/setup-go@v2
+      with:
+        go-version: ${{ matrix.go-version }}
+
+    - name: Checkout code
+      uses: actions/checkout@v2
+
+    - name: Test
+      run: go test ./...

ecs.go

@@ -17,7 +17,7 @@ import (
 var (
 	ecsClient                      = ecs.New(session.New())
 	ec2client                      = ec2.New(session.New())
-	ecsRegExp                      = regexp.MustCompile(`ECS_CLUSTER=(\S*)`)
+	ecsRegExp                      = regexp.MustCompile(`ECS_CLUSTER=([0-9A-Za-z_\-]*)`)
 	ErrMissingUserData             = errors.New("This instance seems not to have UserData")
 	ErrMissingECSClusterInUserData = errors.New("This instance seems not to have EcsCluster definition in UserData")
 	ErrInstanceTerminated          = errors.New("This instance is already terminated")
@@ -175,13 +175,22 @@ func GetClusterNameFromInstanceUserData(ec2Instance string) (string, error) {
 	}
 
 	// Using RegExp to get actual ECS Cluster name from UserData string
-	m := ecsRegExp.FindAllStringSubmatch(string(decodedUserData), -1)
+	val, err := parseECSClusterValue(string(decodedUserData))
+	if err != nil {
+		return "", err
+	}
+
+	return val, nil
+}
+
+// Fetch value of ECS_CLUSTER variable with the regexp
+func parseECSClusterValue(str string) (string, error) {
+	m := ecsRegExp.FindAllStringSubmatch(str, -1)
 	if len(m) == 0 || len(m[0]) < 2 {
-		fmt.Printf("UserData:\n%s", string(decodedUserData))
+		fmt.Printf("UserData:\n%s\n", str)
 		return "", ErrMissingECSClusterInUserData
 	}
 
-	// getting ECS Cluster name which we got from UserData
 	return m[0][1], nil
 }
 

ecs_test.go

@@ -0,0 +1,78 @@
+package drain
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestECSVariable(t *testing.T) {
+	var testcases = []struct {
+		testString    string
+		expectedValue string
+		expectedError error
+	}{
+		{
+			testString:    `echo "ECS_CLUSTER=test-ecs-cluster123" >> /etc/ecs/ecs.config`,
+			expectedValue: "test-ecs-cluster123",
+		},
+		{
+			testString:    `echo 'ECS_CLUSTER=test-ecs-cluster123' >> /etc/ecs/ecs.config`,
+			expectedValue: "test-ecs-cluster123",
+		},
+		{
+			testString:    `echo ECS_CLUSTER=test-ecs-cluster123 >> /etc/ecs/ecs.config`,
+			expectedValue: "test-ecs-cluster123",
+		},
+		{
+			testString: `#!/bin/bash -xe
+			echo ECS_CLUSTER=test-ecs-cluster123 >> /etc/ecs/ecs.config`,
+			expectedValue: "test-ecs-cluster123",
+		},
+		{
+			testString: `
+			write_files:
+			- path: /etc/ecs/ecs.config
+			  append: true
+			  content: |
+				ECS_CLUSTER=my-super-ECS-cluster-22abc
+				ECS_CONTAINER_STOP_TIMEOUT=2m
+`,
+			expectedValue: "my-super-ECS-cluster-22abc",
+		},
+		{
+			testString: `
+			write_files:
+			- path: /etc/ecs/ecs.config
+			  append: true
+			  content: |
+				ECS_CONTAINER_STOP_TIMEOUT=2m
+`,
+			expectedError: ErrMissingECSClusterInUserData,
+		},
+		{
+			testString:    "",
+			expectedError: ErrMissingECSClusterInUserData,
+		},
+		{
+			testString: `
+pip install awscli
+aws configure set default.region ${AWS::Region}
+echo ECS_CLUSTER=my-super-ECS-cluster-22abc >> /etc/ecs/ecs.config
+`,
+			expectedValue: "my-super-ECS-cluster-22abc",
+		},
+	}
+
+	for i, tst := range testcases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			val, err := parseECSClusterValue(tst.testString)
+			if err != tst.expectedError {
+				t.Errorf("Expected %v error, but got %v", tst.expectedError, err)
+			}
+
+			if val != tst.expectedValue {
+				t.Errorf("Expected %q value, but got %q", tst.expectedValue, val)
+			}
+		})
+	}
+}