Support channels in senza images/senza patch

senza/cli.py

@@ -16,8 +16,10 @@
 
 import boto3
 import click
+import senza.stups.taupage as taupage
 import requests
 import yaml
+
 from botocore.exceptions import ClientError
 from clickclick import (Action, FloatRange, OutputFormat, choice, error,
                         fatal_error, info, ok)
@@ -31,7 +33,6 @@
                   get_tag, matches_any, parse_time, resolve_topic_arn,
                   update_stack_from_template, all_stacks_in_final_state)
 from .components import evaluate_template, get_component
-from .components.stups_auto_configuration import find_taupage_image
 from .definitions import AccountArguments
 from .error_handling import HandleExceptions
 from .exceptions import InvalidDefinition, InvalidParameterFile
@@ -1239,10 +1240,11 @@ def images(stack_ref, region, output, field, hide_older_than, show_instances):
     for image in ec2.images.filter(ImageIds=list(instances_by_image.keys())):
         images[image.id] = image
     if not stack_refs:
-        filters = [{'Name': 'name', 'Values': ['*Taupage-*']},
-                   {'Name': 'state', 'Values': ['available']}]
-        for image in ec2.images.filter(Filters=filters):
-            images[image.id] = image
+        for channel in taupage.CHANNELS.values():
+            filters = [{'Name': 'name', 'Values': [channel.ami_wildcard]},
+                       {'Name': 'state', 'Values': ['available']}]
+            for image in ec2.images.filter(Filters=filters):
+                images[image.id] = image
     rows = []
     cutoff = datetime.datetime.now() - datetime.timedelta(days=hide_older_than)
     for image in images.values():
@@ -1409,7 +1411,7 @@ def get_auto_scaling_groups(stack_refs, region):
 @region_option
 @click.option('--image',
               metavar='AMI_ID_OR_LATEST',
-              help='Use specified image (AMI ID or "latest")')
+              help='Use specified image (AMI ID or one of "latest", "staging", "dev")')
 @click.option('--instance-type',
               metavar='INSTANCE_TYPE',
               help='Use specified EC2 instance type')
@@ -1425,8 +1427,11 @@ def patch(stack_ref, region, image, instance_type, user_data):
     region = get_region(region)
     check_credentials(region)
 
-    if image == 'latest':
-        image = find_taupage_image(region).id
+    if image in taupage.CHANNELS:
+        ami = taupage.find_image(region, channel=taupage.CHANNELS[image])
+        if ami is None:
+            fatal_error("No Taupage AMI found for {}".format(image))
+        image = ami.id
 
     properties = {'ImageId': image,
                   'InstanceType': instance_type,

senza/components/stups_auto_configuration.py

@@ -1,47 +1,18 @@
-import boto3
+import senza.stups.taupage as taupage
 
 from senza.components.subnet_auto_configuration import component_subnet_auto_configuration
 from senza.utils import ensure_keys
 
-_TAUPAGE_CHANNELS = {
-    "LatestTaupageImage": "",
-    "LatestTaupageStagingImage": "Staging",
-    "LatestTaupageDevImage": "Dev"
-}
-
-
-def find_taupage_image(region: str, channel_suffix: str = ""):
-    '''Find the latest Taupage AMI, first try private images, fallback to public'''
-    ec2 = boto3.resource('ec2', region)
-    filters = [{'Name': 'name', 'Values': ['*Taupage{}-AMI-*'.format(channel_suffix)]},
-               {'Name': 'is-public', 'Values': ['false']},
-               {'Name': 'state', 'Values': ['available']},
-               {'Name': 'root-device-type', 'Values': ['ebs']}]
-    images = list(ec2.images.filter(Filters=filters))
-    if not images:
-        public_filters = [{'Name': 'name', 'Values': ['*Taupage{}-Public-AMI-*'.format(channel_suffix)]},
-                          {'Name': 'is-public', 'Values': ['true']},
-                          {'Name': 'state', 'Values': ['available']},
-                          {'Name': 'root-device-type', 'Values': ['ebs']}]
-        images = list(ec2.images.filter(Filters=public_filters))
-
-    if not images:
-        if channel_suffix:
-            return None
-        else:
-            # Require at least one image from the stable channel
-            raise Exception('No Taupage AMI found')
-
-    most_recent_image = sorted(images, key=lambda i: i.name)[-1]
-    return most_recent_image
-
 
 def component_stups_auto_configuration(definition, configuration, args, info, force, account_info):
-    for mapping, suffix in _TAUPAGE_CHANNELS.items():
-        most_recent_image = find_taupage_image(args.region, suffix)
+    for channel in taupage.CHANNELS.values():
+        most_recent_image = taupage.find_image(args.region, channel)
         if most_recent_image:
-            configuration = ensure_keys(configuration, "Images", mapping, args.region)
-            configuration["Images"][mapping][args.region] = most_recent_image.id
+            configuration = ensure_keys(configuration, "Images", channel.image_mapping, args.region)
+            configuration["Images"][channel.image_mapping][args.region] = most_recent_image.id
+        elif channel == taupage.DEFAULT_CHANNEL:
+            # Require at least one image from the stable channel
+            raise Exception('No Taupage AMI found')
 
     component_subnet_auto_configuration(definition, configuration, args, info, force, account_info)
 

senza/stups/taupage.py

@@ -0,0 +1,45 @@
+import boto3
+from collections import namedtuple
+
+TaupageChannel = namedtuple("TaupageChannel", ("image_mapping", "ami_wildcard", "public_ami_wildcard"))
+
+
+def _channel(suffix):
+    return TaupageChannel("LatestTaupage{}Image".format(suffix),
+                          "Taupage{}-AMI-*".format(suffix),
+                          "Taupage{}-Public-AMI-*".format(suffix))
+
+
+DEFAULT_CHANNEL = _channel("")
+
+CHANNELS = {
+    "latest": DEFAULT_CHANNEL,
+    "staging": _channel("Staging"),
+    "dev": _channel("Dev")
+}
+
+
+def find_image(region: str, channel: TaupageChannel = None):
+    '''Find the latest Taupage AMI, first try private images, fallback to public'''
+
+    if channel is None:
+        channel = DEFAULT_CHANNEL
+
+    ec2 = boto3.resource('ec2', region)
+    filters = [{'Name': 'name', 'Values': [channel.ami_wildcard]},
+               {'Name': 'is-public', 'Values': ['false']},
+               {'Name': 'state', 'Values': ['available']},
+               {'Name': 'root-device-type', 'Values': ['ebs']}]
+    images = list(ec2.images.filter(Filters=filters))
+    if not images:
+        public_filters = [{'Name': 'name', 'Values': [channel.public_ami_wildcard]},
+                          {'Name': 'is-public', 'Values': ['true']},
+                          {'Name': 'state', 'Values': ['available']},
+                          {'Name': 'root-device-type', 'Values': ['ebs']}]
+        images = list(ec2.images.filter(Filters=public_filters))
+
+    if not images:
+        return None
+
+    most_recent_image = sorted(images, key=lambda i: i.name)[-1]
+    return most_recent_image

tests/test_cli.py

@@ -1551,7 +1551,7 @@ def patch_auto_scaling_group(group, region, properties):
         props.update(properties)
 
     monkeypatch.setattr('boto3.client', MagicMock(return_value=boto3))
-    monkeypatch.setattr('senza.cli.find_taupage_image', MagicMock(return_value=image))
+    monkeypatch.setattr('senza.stups.taupage.find_image', MagicMock(return_value=image))
     monkeypatch.setattr('senza.cli.patch_auto_scaling_group', patch_auto_scaling_group)
     runner = CliRunner()
     result = runner.invoke(cli, ['patch', 'myapp', '1', '--image=latest', '--region=aa-fakeregion-1'],