Scale for elastigroup (#525)

Contributes to #517

senza/aws.py

@@ -301,10 +301,13 @@ def get_stacks(stack_refs: list,
     # stack names that were already yielded to avoid yielding old deleted
     # stacks whose name was reused
     stacks_yielded = set()
+    output_stacks = []
     for stack in stacks:
         if stack['StackName'] not in stacks_yielded or not unique_only:
             stacks_yielded.add(stack['StackName'])
-            yield SenzaStackSummary(stack)
+            output_stacks.append(SenzaStackSummary(stack))
+
+    return output_stacks
 
 
 def matches_any(cf_stack_name: str, stack_refs: list):
@@ -461,7 +464,7 @@ def all_stacks_in_final_state(related_stacks_refs: list, region: str, timeout: O
         while not all_in_final_state and wait_timeout > datetime.datetime.utcnow():
             # assume all stacks are ready
             all_in_final_state = True
-            related_stacks = list(get_stacks(related_stacks_refs, region))
+            related_stacks = get_stacks(related_stacks_refs, region)
 
             if not related_stacks:
                 error("Stack not found!")

senza/cli.py

@@ -25,6 +25,7 @@
                         fatal_error, info, ok)
 from clickclick.console import print_table
 
+from spotinst.components import elastigroup_api
 from .arguments import (GLOBAL_OPTIONS, json_output_option, output_option,
                         parameter_file_option, region_option,
                         stacktrace_visible_option, watch_option,
@@ -118,6 +119,12 @@
 
 SENZA_KMS_PREFIX = 'senza:kms:'
 
+ELASTIGROUP_TYPE = 'Custom::elastigroup'
+
+AUTO_SCALING_GROUP_TYPE = 'AWS::AutoScaling::AutoScalingGroup'
+
+VALID_AUTO_SCALING_GROUPS = [AUTO_SCALING_GROUP_TYPE, ELASTIGROUP_TYPE]
+
 
 def filter_output_columns(output_columns, filter_columns):
     """
@@ -1426,6 +1433,23 @@ def dump(stack_ref, region, output):
         print_json(cfjson, output)
 
 
+def get_auto_scaling_groups_and_elasti_groups(stacks, region):
+    '''
+    Returns data for both AWS Auto Scaling Groups and ElastiGroups
+
+    Note: This method will eventually replace get_auto_scaling_groups when the remaining commands support ElastiGroups
+    '''
+    cf = BotoClientProxy('cloudformation', region)
+    for stack in stacks:
+        resources = cf.describe_stack_resources(StackName=stack.StackName)['StackResources']
+
+        for resource in resources:
+            if resource['ResourceType'] in VALID_AUTO_SCALING_GROUPS:
+                yield {'type': resource['ResourceType'],
+                       'resource_id': resource['PhysicalResourceId'],
+                       'stack_name': resource['StackName']}
+
+
 def get_auto_scaling_groups(stack_refs, region):
     cf = BotoClientProxy('cloudformation', region)
     for stack in get_stacks(stack_refs, region):
@@ -1523,27 +1547,67 @@ def scale(stack_ref, region, desired_capacity, force):
 
     asg = BotoClientProxy('autoscaling', region)
 
-    stack_count = len(list(get_stacks(stack_refs, region)))
+    stacks = get_stacks(stack_refs, region)
+    stack_count = len(stacks)
     if not force and stack_count > 1:
         confirm_str = 'Number of stacks to be scaled - {}. Do you want to continue?'.format(stack_count)
         click.confirm(confirm_str, abort=True)
 
-    for asg_name in get_auto_scaling_groups(stack_refs, region):
-        group = get_auto_scaling_group(asg, asg_name)
-        current_capacity = group['DesiredCapacity']
-        with Action('Scaling {} from {} to {} instances..'.format(
-                asg_name, current_capacity, desired_capacity)) as act:
-            if current_capacity == desired_capacity:
-                act.ok('NO CHANGES')
-            else:
-                kwargs = {}
-                if desired_capacity < group['MinSize']:
-                    kwargs['MinSize'] = desired_capacity
-                if desired_capacity > group['MaxSize']:
-                    kwargs['MaxSize'] = desired_capacity
-                asg.update_auto_scaling_group(AutoScalingGroupName=asg_name,
-                                              DesiredCapacity=desired_capacity,
-                                              **kwargs)
+    for group in get_auto_scaling_groups_and_elasti_groups(stacks, region):
+        if group['type'] == AUTO_SCALING_GROUP_TYPE:
+            scale_auto_scaling_group(asg, group['resource_id'], desired_capacity)
+        elif group['type'] == ELASTIGROUP_TYPE:
+            scale_elastigroup(group['resource_id'], group['stack_name'], desired_capacity, region)
+
+
+def scale_elastigroup(elastigroup_id, stack_name, desired_capacity, region):
+    '''
+    Commands to scale an ElastiGroup
+    '''
+    cf = boto3.client('cloudformation', region)
+    template = cf.get_template(StackName=stack_name)['TemplateBody']
+
+    spotinst_token = template['Mappings']['Senza']['Info']['SpotinstAccessToken']
+    spotinst_account_id = template['Resources']['AppServerConfig']['Properties']['accountId']
+
+    group = elastigroup_api.get_elastigroup(elastigroup_id, spotinst_account_id, spotinst_token)
+    capacity = group['capacity']
+
+    with Action('Scaling ElastiGroup {} (ID: {}) from {} to {} instances..'.format(
+            stack_name, elastigroup_id, capacity['target'], desired_capacity)) as act:
+        if capacity['target'] == desired_capacity:
+            act.ok('NO CHANGES')
+        else:
+            minimum = desired_capacity if capacity['minimum'] > desired_capacity else capacity['minimum']
+            maximum = desired_capacity if capacity['maximum'] < desired_capacity else capacity['maximum']
+
+            elastigroup_api.update_capacity(minimum,
+                                            maximum,
+                                            desired_capacity,
+                                            elastigroup_id,
+                                            spotinst_account_id,
+                                            spotinst_token)
+
+
+def scale_auto_scaling_group(asg, asg_name, desired_capacity):
+    '''
+    Commands to scale an AWS Auto Scaling Group
+    '''
+    group = get_auto_scaling_group(asg, asg_name)
+    current_capacity = group['DesiredCapacity']
+    with Action('Scaling {} from {} to {} instances..'.format(
+            asg_name, current_capacity, desired_capacity)) as act:
+        if current_capacity == desired_capacity:
+            act.ok('NO CHANGES')
+        else:
+            kwargs = {}
+            if desired_capacity < group['MinSize']:
+                kwargs['MinSize'] = desired_capacity
+            if desired_capacity > group['MaxSize']:
+                kwargs['MaxSize'] = desired_capacity
+            asg.update_auto_scaling_group(AutoScalingGroupName=asg_name,
+                                          DesiredCapacity=desired_capacity,
+                                          **kwargs)
 
 
 def failure_event(event: dict):
@@ -1586,7 +1650,7 @@ def wait(stack_ref, region, deletion, timeout, interval):
         stacks_in_progress = set()
         successful_actions = set()
 
-        stacks_found = list(get_stacks(stack_refs, region, all=True, unique_only=True))
+        stacks_found = get_stacks(stack_refs, region, all=True, unique_only=True)
 
         if not stacks_found:
             raise click.UsageError('No matching stack for "{}" found'.format(' '.join(stack_ref)))

spotinst/components/elastigroup_api.py

@@ -0,0 +1,60 @@
+'''
+Wrapper methods for ElastiGroup's API
+'''
+import requests
+import json
+
+
+SPOTINST_API_URL = 'https://api.spotinst.io'
+
+
+def update_capacity(minimum, maximum, target, elastigroup_id, spotinst_account_id, spotinst_token):
+    '''
+    Updates the capacity (number of instances) for an ElastiGroup by calling the SpotInst API.
+    Returns the updated description of the ElastiGroup as a dict.
+    Exceptions will be thrown for HTTP errors.
+
+    For more details see https://api.spotinst.com/elastigroup/amazon-web-services/update/
+    '''
+    headers = {
+        "Authorization": "Bearer {}".format(spotinst_token),
+        "Content-Type": "application/json"
+    }
+
+    new_capacity = {
+        'target': target,
+        'minimum': minimum,
+        'maximum': maximum
+    }
+
+    body = {'group': {'capacity': new_capacity}}
+    response = requests.put(
+        '{}/aws/ec2/group/{}?accountId={}'.format(SPOTINST_API_URL, elastigroup_id, spotinst_account_id),
+        headers=headers, timeout=10, data=json.dumps(body))
+    response.raise_for_status()
+    data = response.json()
+    groups = data.get("response", {}).get("items", [])
+
+    return groups[0]
+
+
+def get_elastigroup(elastigroup_id, spotinst_account_id, spotinst_token):
+    '''
+    Returns the description of an ElastiGroup as a dict.
+    Exceptions will be thrown for HTTP errors.
+
+    For more details see https://api.spotinst.com/elastigroup/amazon-web-services/list-group/
+    '''
+    headers = {
+        "Authorization": "Bearer {}".format(spotinst_token),
+        "Content-Type": "application/json"
+    }
+
+    response = requests.get(
+        'https://api.spotinst.io/aws/ec2/group/{}?accountId={}'.format(elastigroup_id, spotinst_account_id),
+        headers=headers, timeout=5)
+    response.raise_for_status()
+    data = response.json()
+    groups = data.get("response", {}).get("items", [])
+
+    return groups[0]

tests/test_cli.py

@@ -5,7 +5,6 @@
 from contextlib import contextmanager
 
 from typing import List, Dict
-from typing import Optional
 from unittest.mock import MagicMock, mock_open
 
 import botocore.exceptions
@@ -1645,7 +1644,8 @@ def test_scale(monkeypatch):
                                                           'CreationTime': '2016-06-14'}]}
     boto3.describe_stack_resources.return_value = {'StackResources':
                                                        [{'ResourceType': 'AWS::AutoScaling::AutoScalingGroup',
-                                                         'PhysicalResourceId': 'myasg'}]}
+                                                         'PhysicalResourceId': 'myasg',
+                                                         'StackName': 'myapp-1'}]}
     # NOTE: we are using invalid MinSize (< capacity) here to get one more line covered ;-)
     group = {'AutoScalingGroupName': 'myasg', 'DesiredCapacity': 1, 'MinSize': 3, 'MaxSize': 1}
     boto3.describe_auto_scaling_groups.return_value = {'AutoScalingGroups': [group]}
@@ -1655,6 +1655,68 @@ def test_scale(monkeypatch):
                            catch_exceptions=False)
     assert 'Scaling myasg from 1 to 2 instances' in result.output
 
+
+def test_scale_elastigroup(monkeypatch):
+    spotinst_account_id = 'fakeactid'
+    elastigroup_id = 'myelasti'
+    boto3 = MagicMock()
+    boto3.list_stacks.return_value = {'StackSummaries': [{'StackName': 'myapp-1',
+                                                          'CreationTime': '2016-06-14'}]}
+    boto3.describe_stack_resources.return_value = {'StackResources':
+                                                       [{'ResourceType': 'Custom::elastigroup',
+                                                         'PhysicalResourceId': elastigroup_id,
+                                                         'StackName': 'myapp-1'}]}
+    boto3.get_template.return_value = {
+        'TemplateBody': {
+            'Mappings': {
+                'Senza': {
+                    'Info': {
+                        'SpotinstAccessToken': 'faketoken'
+                    }
+                }
+            },
+            'Resources': {
+                'AppServerConfig': {
+                    'Properties': {
+                        'accountId': spotinst_account_id
+                    }
+                }
+            }
+        }
+    }
+
+    group = {
+        'capacity': {
+            'minimum': 1,
+            'maximum': 2,
+            'target': 1,
+            'unit': 'instance'
+        }
+    }
+    get_elastigroup = MagicMock()
+    get_elastigroup.return_value = group
+
+    update = {
+        'capacity': {
+            'minimum': 1,
+            'maximum': 3,
+            'target': 3,
+            'unit': 'instance'
+        }
+    }
+    update_capacity = MagicMock()
+    update_capacity.return_value = update
+
+    monkeypatch.setattr('boto3.client', MagicMock(return_value=boto3))
+    monkeypatch.setattr('spotinst.components.elastigroup_api.get_elastigroup', get_elastigroup)
+    monkeypatch.setattr('spotinst.components.elastigroup_api.update_capacity', update_capacity)
+
+    runner = CliRunner()
+    result = runner.invoke(cli, ['scale', 'myapp', '1', '3', '--region=aa-fakeregion-1'],
+                           catch_exceptions=False)
+    assert 'Scaling ElastiGroup myapp-1 (ID: myelasti) from 1 to 3 instances' in result.output
+
+
 def test_scale_with_confirm(monkeypatch):
     boto3 = MagicMock()
     boto3.list_stacks.return_value = {'StackSummaries': [
@@ -1683,7 +1745,8 @@ def test_scale_with_force_confirm(monkeypatch):
     ]}
     boto3.describe_stack_resources.return_value = {'StackResources':
                                                        [{'ResourceType': 'AWS::AutoScaling::AutoScalingGroup',
-                                                         'PhysicalResourceId': 'myasg'}]}
+                                                         'PhysicalResourceId': 'myasg',
+                                                         'StackName': 'myapp-1'}]}
     group = {'AutoScalingGroupName': 'myasg', 'DesiredCapacity': 1, 'MinSize': 3, 'MaxSize': 1}
     boto3.describe_auto_scaling_groups.return_value = {'AutoScalingGroups': [group]}
     monkeypatch.setattr('boto3.client', MagicMock(return_value=boto3))
@@ -1693,7 +1756,6 @@ def test_scale_with_force_confirm(monkeypatch):
     assert 'Scaling myasg from 1 to 2 instances' in result.output
 
 
-
 def test_wait(monkeypatch):
     cf = MagicMock()
     stack1 = {'StackName': 'test-1',

tests/test_elastigroup_api.py

@@ -0,0 +1,59 @@
+import responses
+from spotinst.components.elastigroup_api import update_capacity, get_elastigroup, SPOTINST_API_URL
+
+
+def test_update_capacity(monkeypatch):
+    update = {
+        'response': {
+            'items': [{
+                'id': 'sig-xfy',
+                'name': 'my-app-1',
+                'capacity': {
+                    'minimum': 1,
+                    'maximum': 3,
+                    'target': 3,
+                    'unit': 'instance'
+                }
+            }]
+        }
+    }
+
+    elastigroup_id = 'sig-xfy'
+    spotinst_account_id = 'act-zwk'
+    spotinst_token = 'fake-token'
+    with responses.RequestsMock() as rsps:
+        rsps.add(rsps.PUT,
+                 '{}/aws/ec2/group/{}?accountId={}'.format(SPOTINST_API_URL, elastigroup_id, spotinst_account_id),
+                 status=200,
+                 json=update)
+
+        update_response = update_capacity(1, 3, 3, elastigroup_id, spotinst_account_id, spotinst_token)
+        assert update_response['id'] == elastigroup_id
+        assert update_response['name'] == 'my-app-1'
+        assert update_response['capacity']['minimum'] == 1
+        assert update_response['capacity']['maximum'] == 3
+        assert update_response['capacity']['target'] == 3
+
+
+def test_get_elastigroup(monkeypatch):
+    group = {
+        'response': {
+            'items': [{
+                'id': 'sig-xfy',
+                'name': 'my-app-1',
+                'region': 'eu-central-1',
+            }]
+        }
+    }
+
+    elastigroup_id = 'sig-xfy'
+    spotinst_account_id = 'act-zwk'
+    spotinst_token = 'fake-token'
+    with responses.RequestsMock() as rsps:
+        rsps.add(rsps.GET, 'https://api.spotinst.io/aws/ec2/group/{}?accountId={}'.format(elastigroup_id, spotinst_account_id),
+                 status=200,
+                 json=group)
+
+        group = get_elastigroup(elastigroup_id, spotinst_account_id, spotinst_token)
+        assert group['id'] == elastigroup_id
+        assert group['name'] == 'my-app-1'