From 72d43e340ab6f603e0655216b7f0447954320d2c Mon Sep 17 00:00:00 2001
From: Luis <lmineiro@gmail.com>
Date: Tue, 18 Sep 2018 12:56:54 +0200
Subject: [PATCH] added support for the TargetGroupARNs property in
 Elastigroups (#536)

- adds support for the `TargetGroupARNs` property. When present, they will be set as the [Elastigroup's `loadBalancers`](https://api.spotinst.com/spotinst-api/elastigroup/amazon-web-services/create/).
---
 senza/components/elastigroup.py | 24 ++++++++-----
 tests/test_elastigroup.py       | 60 +++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 8 deletions(-)

diff --git a/senza/components/elastigroup.py b/senza/components/elastigroup.py
index 301359a0..0c70d33f 100644
--- a/senza/components/elastigroup.py
+++ b/senza/components/elastigroup.py
@@ -351,17 +351,25 @@ def extract_load_balancer_name(configuration, elastigroup_config: dict):
         if "ElasticLoadBalancerV2" in configuration:
             health_check_type = "TARGET_GROUP"
             load_balancer_refs = configuration.pop("ElasticLoadBalancerV2")
-            if isinstance(load_balancer_refs, str):
-                load_balancers.append({
-                    "arn": {"Ref": load_balancer_refs + 'TargetGroup'},
-                    "type": "TARGET_GROUP"
-                })
-            elif isinstance(load_balancer_refs, list):
-                for load_balancer_ref in load_balancer_refs:
+            custom_target_groups = configuration.pop("TargetGroupARNs", None)
+            if custom_target_groups:
+                for custom_target_group in custom_target_groups:
+                    load_balancers.append({
+                        "arn": custom_target_group,
+                        "type": "TARGET_GROUP"
+                    })
+            else:
+                if isinstance(load_balancer_refs, str):
                     load_balancers.append({
-                        "arn": {"Ref": load_balancer_ref + "TargetGroup"},
+                        "arn": {"Ref": load_balancer_refs + 'TargetGroup'},
                         "type": "TARGET_GROUP"
                     })
+                elif isinstance(load_balancer_refs, list):
+                    for load_balancer_ref in load_balancer_refs:
+                        load_balancers.append({
+                            "arn": {"Ref": load_balancer_ref + "TargetGroup"},
+                            "type": "TARGET_GROUP"
+                        })
 
         if len(load_balancers) > 0:
             launch_spec_config["loadBalancersConfig"] = {"loadBalancers": load_balancers}
diff --git a/tests/test_elastigroup.py b/tests/test_elastigroup.py
index b98621e6..5648a830 100644
--- a/tests/test_elastigroup.py
+++ b/tests/test_elastigroup.py
@@ -508,6 +508,66 @@ def test_load_balancers():
         assert test_case["expected_config"] == got
 
 
+def test_multiple_target_groups():
+    test_cases = [
+        {  # multiple target groups in raw ARN form, ignore ALB TG
+            "input": {"ElasticLoadBalancerV2": "foo", "TargetGroupARNs": ["bar", "baz"]},
+            "given_config": {},
+            "expected_config": {"compute": {"launchSpecification": {
+                "loadBalancersConfig": {
+                    "loadBalancers": [
+                        {"arn": "bar", "type": "TARGET_GROUP"},
+                        {"arn": "baz", "type": "TARGET_GROUP"},
+                    ],
+                },
+                "healthCheckType": "TARGET_GROUP",
+                "healthCheckGracePeriod": 300,
+            }}},
+        },
+        {  # multiple target groups with Ref, ignore ALB TG
+            "input": {"ElasticLoadBalancerV2": "foo", "TargetGroupARNs": [{"Ref": "bar"}, {"Ref": "baz"}]},
+            "given_config": {},
+            "expected_config": {"compute": {"launchSpecification": {
+                "loadBalancersConfig": {
+                    "loadBalancers": [
+                        {"arn": {"Ref": "bar"}, "type": "TARGET_GROUP"},
+                        {"arn": {"Ref": "baz"}, "type": "TARGET_GROUP"},
+                    ],
+                },
+                "healthCheckType": "TARGET_GROUP",
+                "healthCheckGracePeriod": 300,
+            }}},
+        },
+        {  # ignore Taupage target groups, leave Elatigroup untouched
+            "input": {"ElasticLoadBalancerV2": "foo", "TargetGroupARNs": [{"Ref": "bar"}, {"Ref": "baz"}]},
+            "given_config": {"compute": {
+                "launchSpecification": {
+                    "loadBalancersConfig": {
+                        "loadBalancers": [
+                            {"arn": "givenTargetGroup", "type": "TARGET_GROUP"},
+                        ],
+                    },
+                    "healthCheckType": "TARGET_GROUP",
+                    "healthCheckGracePeriod": 300,
+                }
+            }},
+            "expected_config": {"compute": {"launchSpecification": {
+                "loadBalancersConfig": {
+                    "loadBalancers": [
+                        {"arn": "givenTargetGroup", "type": "TARGET_GROUP"},
+                    ],
+                },
+                "healthCheckType": "TARGET_GROUP",
+                "healthCheckGracePeriod": 300,
+            }}},
+        },
+    ]
+    for test_case in test_cases:
+        got = test_case["given_config"]
+        extract_load_balancer_name(test_case["input"], got)
+        assert test_case["expected_config"] == got
+
+
 def test_public_ips():
     test_cases = [
         {  # default behavior - no public IPs, leave untouched