Merge pull request #128 from theserverlessway/feature/add-only-missin…

…g-stack-set-instances

Tests for adding only missing stack set instances

formica/cli.py

@@ -29,6 +29,7 @@
     'profile': str,
     'capabilities': list,
     'vars': dict,
+    'timeout': int,
     'administration_role_arn': str,
     'administration_role_name': str,
     'execution_role_name': str,
@@ -133,6 +134,7 @@ def main(cli_args):
     add_aws_arguments(deploy_parser)
     add_stack_argument(deploy_parser)
     add_config_file_argument(deploy_parser)
+    add_timeout_parameter(deploy_parser)
     deploy_parser.set_defaults(func=deploy)
 
     # Cancel Command Arguments
@@ -289,6 +291,7 @@ def stack_set_parser(parser):
     add_config_file_argument(add_instances_parser)
     add_stack_set_main_auto_regions_accounts(add_instances_parser)
     add_stack_set_operation_preferences(add_instances_parser)
+    add_yes_parameter(add_instances_parser)
     add_instances_parser.set_defaults(func=stack_set.add_stack_set_instances)
 
     # Remove Instances
@@ -301,6 +304,7 @@ def stack_set_parser(parser):
     add_config_file_argument(remove_instances_parser)
     add_stack_set_main_auto_regions_accounts(remove_instances_parser)
     add_stack_set_operation_preferences(remove_instances_parser)
+    add_yes_parameter(remove_instances_parser)
     remove_instances_parser.set_defaults(func=stack_set.remove_stack_set_instances)
 
     # Diff
@@ -435,6 +439,10 @@ def add_yes_parameter(parser):
     parser.add_argument('--yes', '-y', help='Answer all input questions with yes', action='store_true')
 
 
+def add_timeout_parameter(parser):
+    parser.add_argument('--timeout', help='Set the Timeout in minutes before the Update is canceled', type=int)
+
+
 def template(args):
     from .loader import Loader
     import yaml
@@ -530,15 +538,18 @@ def stack_wait_handler(args):
         stack_id = client.describe_stacks(StackName=args.stack)['Stacks'][0]['StackId']
         last_event = client.describe_stack_events(StackName=args.stack)['StackEvents'][0]['EventId']
         function(args, client)
-        StackWaiter(stack_id, client).wait(last_event)
+        options = {}
+        if vars(args).get('timeout'):
+            options['timeout'] = args.timeout
+        StackWaiter(stack_id, client, **options).wait(last_event)
 
     return stack_wait_handler
 
 
 @requires_stack
 @wait_for_stack
 def deploy(args, client):
-    logger.info('Deploying StackSet to {}'.format(args.stack))
+    logger.info('Deploying Stack to {}'.format(args.stack))
     client.execute_change_set(ChangeSetName=(CHANGE_SET_FORMAT.format(stack=args.stack)), StackName=args.stack)
 
 

formica/stack_set.py

@@ -5,6 +5,7 @@
 from .aws import AWS
 from .helper import collect_vars, main_account_id, aws_accounts, aws_regions
 from .diff import compare_stack_set
+from texttable import Texttable
 
 logger = logging.getLogger(__name__)
 
@@ -36,11 +37,15 @@ def validate_stack_set(args):
     return validate_stack_set
 
 
+def ack(message):
+    return input('{}: [y/yes]: '.format(message)).lower() in ['y', 'yes']
+
+
 @requires_stack_set
 def update_stack_set(args):
     compare_stack_set(stack=args.stack_set, vars=collect_vars(args), parameters=args.parameters, tags=args.tags,
                       main_account_parameter=args.main_account_parameter)
-    if args.yes or input('Do you want to update the StackSet with above changes: [y/yes]: ').lower() in ['y', 'yes']:
+    if args.yes or ack('Do you want to update the StackSet with above changes'):
         __manage_stack_set(args=args, create=False)
     else:
         logger.info('StackSet Update canceled')
@@ -59,31 +64,77 @@ def remove_stack_set(args):
     logger.info('Removed StackSet with name {}'.format(args.stack_set))
 
 
+def accounts_table(accounts_map):
+    table = Texttable(max_width=150)
+    table.set_cols_dtype(['t', 't'])
+    table.add_rows([['Account', 'Regions']])
+
+    for account, regions in accounts_map.items():
+        table.add_row([account, ', '.join(regions)])
+
+    logger.info(table.draw() + "\n")
+
+
 @requires_stack_set
 @requires_accounts_regions
 def add_stack_set_instances(args):
     client = AWS.current_session().client('cloudformation')
-    preferences = operation_preferences(args)
-    result = client.create_stack_instances(StackSetName=args.stack_set,
-                                           Accounts=accounts(args),
-                                           Regions=regions(args),
-                                           **preferences)
-    logger.info('Adding StackSet Instances for StackSet {}'.format(args.stack_set))
-    wait_for_stack_set_operation(args.stack_set, result['OperationId'])
+    paginator = client.get_paginator('list_stack_instances')
+    deployed = [{'Account': stack['Account'], 'Region': stack['Region']} for page in
+                paginator.paginate(StackSetName=args.stack_set) for stack in page['Summaries']]
+
+    expected_instances = [{'Account': account, 'Region': region} for account in accounts(args) for region in
+                          regions(args)]
+
+    new_instances = [i for i in expected_instances if i not in deployed]
+    if new_instances:
+        new_accounts = sorted(list(set([i['Account'] for i in new_instances])))
+        new_regions = sorted(list(set([i['Region'] for i in new_instances])))
+
+        account_to_region = {a: set([i['Region'] for i in new_instances if i['Account'] == a]) for a in new_accounts}
+
+        if len(new_instances) == len(expected_instances) or len(new_accounts) == 1 or len(new_regions) == 1:
+            targets = [(new_accounts, new_regions)]
+        else:
+            targets = [([i['Account'] for i in new_instances if i['Region'] == region], [region]) for region in
+                       new_regions]
+
+        logger.info('Adding new StackSet Instances:')
+        accounts_table(account_to_region)
+        if args.yes or ack('Do you want to add these StackSet Instances:'):
+            for target in targets:
+                preferences = operation_preferences(args)
+                result = client.create_stack_instances(StackSetName=args.stack_set,
+                                                       Accounts=target[0],
+                                                       Regions=target[1],
+                                                       **preferences)
+                wait_for_stack_set_operation(args.stack_set, result['OperationId'])
+        else:
+            logger.info('Adding StackSet Instances canceled')
+            sys.exit(1)
+    else:
+        logger.info('All StackSet Instances are deployed')
 
 
 @requires_stack_set
 @requires_accounts_regions
 def remove_stack_set_instances(args):
     client = AWS.current_session().client('cloudformation')
     preferences = operation_preferences(args)
-    result = client.delete_stack_instances(StackSetName=args.stack_set,
-                                           Accounts=accounts(args),
-                                           Regions=regions(args),
-                                           RetainStacks=args.retain,
-                                           **preferences)
+    acc = accounts(args)
+    reg = regions(args)
     logger.info('Removing StackSet Instances for StackSet {}'.format(args.stack_set))
-    wait_for_stack_set_operation(args.stack_set, result['OperationId'])
+    accounts_table({a: reg for a in acc})
+    if args.yes or ack('Do you want to remove these StackSet Instances'):
+        result = client.delete_stack_instances(StackSetName=args.stack_set,
+                                               Accounts=acc,
+                                               Regions=reg,
+                                               RetainStacks=args.retain,
+                                               **preferences)
+        wait_for_stack_set_operation(args.stack_set, result['OperationId'])
+    else:
+        logger.info('Removing StackSet Instances canceled')
+        sys.exit(1)
 
 
 @requires_stack_set

formica/stack_waiter.py

@@ -1,5 +1,6 @@
 import sys
 import time
+from datetime import datetime
 
 import logging
 from texttable import Texttable
@@ -12,20 +13,24 @@
 FAILED_STATES = ['CREATE_FAILED', 'DELETE_FAILED', 'ROLLBACK_FAILED', 'ROLLBACK_COMPLETE', 'UPDATE_FAILED',
                  'UPDATE_ROLLBACK_FAILED', 'UPDATE_ROLLBACK_COMPLETE']
 
-
 logger = logging.getLogger(__name__)
 
+SLEEP_TIME = 5
+
 
 class StackWaiter:
-    def __init__(self, stack, client):
+    def __init__(self, stack, client, timeout=0):
         self.stack = stack
         self.client = client
+        self.timeout = timeout
 
     def wait(self, last_event):
         self.print_header()
         finished = False
+        canceled = False
+        start = datetime.now()
         while not finished:
-            time.sleep(5)
+            time.sleep(SLEEP_TIME)
             stack_events = self.client.describe_stack_events(StackName=self.stack)['StackEvents']
             index = next((i for i, v in enumerate(stack_events) if v['EventId'] == last_event))
             last_event = stack_events[0]['EventId']
@@ -39,13 +44,19 @@ def wait(self, last_event):
             elif stack_status in FAILED_STATES:
                 logger.info("Stack Change Failed: {}".format(stack_status))
                 sys.exit(1)
+            elif not canceled and self.timeout > 0 and (datetime.now() - start).seconds > (self.timeout * 60):
+                logger.info("Timeout of {} minute(s) reached. Canceling Update.".format(self.timeout))
+                canceled = True
+                self.client.cancel_update_stack(StackName=self.stack)
 
     def __create_table(self):
         table = Texttable()
         table.set_cols_width(TABLE_COLUMN_SIZE)
         return table
 
     def print_header(self):
+        if self.timeout > 0:
+            logger.info('Timeout set to {} minute(s)'.format(self.timeout))
         table = self.__create_table()
         table.add_rows([EVENT_TABLE_HEADERS])
         table.set_deco(Texttable.BORDER | Texttable.VLINES)

tests/unit/test_deploy.py

@@ -56,3 +56,14 @@ def test_executes_change_set_and_waits(session, stack_waiter):
     cf_client_mock.execute_change_set.assert_called_with(ChangeSetName=CHANGESETNAME, StackName=STACK)
     stack_waiter.assert_called_with(STACK_ID, cf_client_mock)
     stack_waiter.return_value.wait.assert_called_with(EVENT_ID)
+
+
+def test_executes_change_set_with_timeout(session, stack_waiter):
+    cf_client_mock = Mock()
+    session.return_value.client.return_value = cf_client_mock
+    cf_client_mock.describe_stack_events.return_value = {'StackEvents': [{'EventId': EVENT_ID}]}
+    cf_client_mock.describe_stacks.return_value = {'Stacks': [{'StackId': STACK_ID}]}
+
+    cli.main(['deploy', '--stack', STACK, '--profile', PROFILE, '--region', REGION, '--timeout', '15'])
+    stack_waiter.assert_called_with(STACK_ID, cf_client_mock, timeout=15)
+    stack_waiter.return_value.wait.assert_called_with(EVENT_ID)

tests/unit/test_stack_set.py

@@ -26,6 +26,7 @@ def client(session, mocker):
     client_mock.create_stack_instances.return_value = {'OperationId': OPERATION_ID}
     client_mock.delete_stack_instances.return_value = {'OperationId': OPERATION_ID}
     client_mock.update_stack_set.return_value = {'OperationId': OPERATION_ID}
+    client_mock.get_paginator.return_value.paginate.return_value = []
     return client_mock
 
 
@@ -286,7 +287,7 @@ def test_update_stack_set_with_all_subaccounts(client, logger, loader, input, co
     )
 
 
-def test_add_stack_set_instances(client, loader, wait):
+def test_add_stack_set_instances(client, loader, wait, input):
     cli.main([
         'stack-set',
         'add-instances',
@@ -308,7 +309,7 @@ def test_stack_set_accounts_converts_to_string():
     assert stack_set.accounts({'accounts': ['123', 12345, 54321]}) == ['123', '12345', '54321']
 
 
-def test_add_all_stack_set_instances(client, loader, wait):
+def test_add_all_stack_set_instances(client, loader, wait, input):
     client.list_accounts.return_value = ACCOUNTS
     client.get_caller_identity.return_value = {'Account': '5678'}
     client.describe_regions.return_value = EC2_REGIONS
@@ -327,7 +328,7 @@ def test_add_all_stack_set_instances(client, loader, wait):
     )
 
 
-def test_add_stack_set_instances_with_operation_preferences(client, loader, wait):
+def test_add_stack_set_instances_with_operation_preferences(client, loader, wait, input):
     cli.main([
         'stack-set',
         'add-instances',
@@ -349,7 +350,73 @@ def test_add_stack_set_instances_with_operation_preferences(client, loader, wait
     )
 
 
-def test_remove_stack_set_instances(client, loader, wait):
+def test_add_only_missing_stack_set_instances(client, loader, wait, input, mocker):
+    client.get_paginator.return_value.paginate.return_value = [
+        {'Summaries': [{'Account': '123456789', "Region": 'eu-central-1'}]}]
+    cli.main([
+        'stack-set',
+        'add-instances',
+        '--accounts', '123456789', '987654321',
+        '--regions', 'eu-central-1', 'eu-west-1',
+        '--stack-set', STACK
+    ])
+
+    assert client.create_stack_instances.mock_calls == [
+        mocker.call(
+            StackSetName=STACK,
+            Accounts=['987654321'],
+            Regions=['eu-central-1']
+        ),
+        mocker.call(
+            StackSetName=STACK,
+            Accounts=['123456789', '987654321'],
+            Regions=['eu-west-1']
+        )
+    ]
+
+
+def test_add_only_missing_stack_set_instances_with_one_call_if_possible(client, loader, wait, input, mocker):
+    client.get_paginator.return_value.paginate.return_value = [
+        {'Summaries': [{'Account': '123456789', "Region": 'eu-central-1'},
+                       {'Account': '987654321', "Region": 'eu-central-1'}]}]
+    cli.main([
+        'stack-set',
+        'add-instances',
+        '--accounts', '123456789', '987654321',
+        '--regions', 'eu-central-1', 'eu-west-1',
+        '--stack-set', STACK
+    ])
+
+    assert client.create_stack_instances.mock_calls == [
+        mocker.call(
+            StackSetName=STACK,
+            Accounts=['123456789', '987654321'],
+            Regions=['eu-west-1']
+        )
+    ]
+
+    client.get_paginator.return_value.paginate.return_value = [
+        {'Summaries': [{'Account': '123456789', "Region": 'eu-central-1'},
+                       {'Account': '123456789', "Region": 'eu-west-1'}]}]
+
+    cli.main([
+        'stack-set',
+        'add-instances',
+        '--accounts', '123456789', '987654321',
+        '--regions', 'eu-central-1', 'eu-west-1',
+        '--stack-set', STACK
+    ])
+
+    assert len(client.create_stack_instances.mock_calls) == 2
+
+    client.create_stack_instances.assert_called_with(
+        StackSetName=STACK,
+        Accounts=['987654321'],
+        Regions=['eu-central-1', 'eu-west-1']
+    )
+
+
+def test_remove_stack_set_instances(client, loader, wait, input):
     cli.main([
         'stack-set',
         'remove-instances',
@@ -377,7 +444,7 @@ def test_remove_stack_set_instances(client, loader, wait):
     wait.assert_called_with(STACK, OPERATION_ID)
 
 
-def test_remove_all_stack_set_instances(client, loader, wait):
+def test_remove_all_stack_set_instances(client, loader, wait, input):
     client.list_accounts.return_value = ACCOUNTS
     client.get_caller_identity.return_value = {'Account': '5678'}
     client.describe_regions.return_value = EC2_REGIONS