Refactor ec2_snapshot to make it more testable

lib/ansible/modules/cloud/amazon/ec2_snapshot.py

@@ -74,9 +74,9 @@ options:
       - snapshot id to remove
     required: false
     version_added: "1.9"
-  snapshot_max_age:
+  last_snapshot_min_age:
     description:
-      - If the volume's most recent snapshot has started less than `snapshot_max_age' minutes ago, a new snapshot will not be created.
+      - If the volume's most recent snapshot has started less than `last_snapshot_min_age' minutes ago, a new snapshot will not be created.
     required: false
     default: 0
     version_added: "1.9"
@@ -115,10 +115,11 @@ EXAMPLES = '''
 - local_action:
     module: ec2_snapshot
     volume_id: vol-abcdef12
-    snapshot_max_age: 60
+    last_snapshot_min_age: 60
 '''
 
 import time
+import datetime
 
 try:
     import boto.ec2
@@ -127,7 +128,128 @@ except ImportError:
     HAS_BOTO = False
 
 
-def main():
+# Find the most recent snapshot
+def _get_snapshot_starttime(snap):
+    return datetime.datetime.strptime(snap.start_time, '%Y-%m-%dT%H:%M:%S.000Z')
+
+
+def _get_most_recent_snapshot(snapshots, max_snapshot_age_secs=None, now=None):
+    """
+    Gets the most recently created snapshot and optionally filters the result
+    if the snapshot is too old
+    :param snapshots: list of snapshots to search
+    :param max_snapshot_age_secs: filter the result if its older than this
+    :param now: simulate time -- used for unit testing
+    :return:
+    """
+    if len(snapshots) == 0:
+        return None
+
+    if not now:
+        now = datetime.datetime.utcnow()
+
+    youngest_snapshot = min(snapshots, key=_get_snapshot_starttime)
+
+    # See if the snapshot is younger that the given max age
+    snapshot_start = datetime.datetime.strptime(youngest_snapshot.start_time, '%Y-%m-%dT%H:%M:%S.000Z')
+    snapshot_age = now - snapshot_start
+
+    if max_snapshot_age_secs is not None:
+        if snapshot_age.total_seconds() > max_snapshot_age_secs:
+            return None
+
+    return youngest_snapshot
+
+
+def _create_with_wait(snapshot, wait_timeout_secs, sleep_func=time.sleep):
+    """
+    Wait for the snapshot to be created
+    :param snapshot:
+    :param wait_timeout_secs: fail this step after this many seconds
+    :param sleep_func:
+    :return:
+    """
+    time_waited = 0
+    snapshot.update()
+    while snapshot.status != 'completed':
+        sleep_func(3)
+        snapshot.update()
+        time_waited += 3
+        if wait_timeout_secs and time_waited > wait_timeout_secs:
+            return False
+    return True
+
+
+def create_snapshot(module, ec2, state=None, description=None, wait=None,
+                    wait_timeout=None, volume_id=None, instance_id=None,
+                    snapshot_id=None, device_name=None, snapshot_tags=None,
+                    last_snapshot_min_age=None):
+    snapshot = None
+    changed = False
+
+    required = [volume_id, snapshot_id, instance_id]
+    if required.count(None) != len(required) - 1: # only 1 must be set
+        module.fail_json(msg='One and only one of volume_id or instance_id or snapshot_id must be specified')
+    if instance_id and not device_name or device_name and not instance_id:
+        module.fail_json(msg='Instance ID and device name must both be specified')
+
+    if instance_id:
+        try:
+            volumes = ec2.get_all_volumes(filters={'attachment.instance-id': instance_id, 'attachment.device': device_name})
+        except boto.exception.BotoServerError, e:
+            module.fail_json(msg = "%s: %s" % (e.error_code, e.error_message))
+
+        if not volumes:
+            module.fail_json(msg="Could not find volume with name %s attached to instance %s" % (device_name, instance_id))
+
+        volume_id = volumes[0].id
+
+    if state == 'absent':
+        if not snapshot_id:
+            module.fail_json(msg = 'snapshot_id must be set when state is absent')
+        try:
+            ec2.delete_snapshot(snapshot_id)
+        except boto.exception.BotoServerError, e:
+            # exception is raised if snapshot does not exist
+            if e.error_code == 'InvalidSnapshot.NotFound':
+                module.exit_json(changed=False)
+            else:
+                module.fail_json(msg = "%s: %s" % (e.error_code, e.error_message))
+
+        # successful delete
+        module.exit_json(changed=True)
+
+    if last_snapshot_min_age > 0:
+        try:
+            current_snapshots = ec2.get_all_snapshots(filters={'volume_id': volume_id})
+        except boto.exception.BotoServerError, e:
+            module.fail_json(msg="%s: %s" % (e.error_code, e.error_message))
+
+        last_snapshot_min_age = last_snapshot_min_age * 60 # Convert to seconds
+        snapshot = _get_most_recent_snapshot(current_snapshots,
+                                             max_snapshot_age_secs=last_snapshot_min_age)
+    try:
+        # Create a new snapshot if we didn't find an existing one to use
+        if snapshot is None:
+            snapshot = ec2.create_snapshot(volume_id, description=description)
+            changed = True
+        if wait:
+            if not _create_with_wait(snapshot, wait_timeout):
+                module.fail_json(msg='Timed out while creating snapshot.')
+        if snapshot_tags:
+            for k, v in snapshot_tags.items():
+                snapshot.add_tag(k, v)
+    except boto.exception.BotoServerError, e:
+        module.fail_json(msg="%s: %s" % (e.error_code, e.error_message))
+
+    module.exit_json(changed=changed,
+                     snapshot_id=snapshot.id,
+                     volume_id=snapshot.volume_id,
+                     volume_size=snapshot.volume_size,
+                     tags=snapshot.tags.copy())
+
+
+def create_snapshot_ansible_module():
     argument_spec = ec2_argument_spec()
     argument_spec.update(
         dict(
@@ -138,12 +260,17 @@ def main():
             device_name = dict(),
             wait = dict(type='bool', default=True),
             wait_timeout = dict(type='int', default=0),
-            snapshot_max_age = dict(type='int', default=0),
+            last_snapshot_min_age = dict(type='int', default=0),
             snapshot_tags = dict(type='dict', default=dict()),
             state = dict(choices=['absent','present'], default='present'),
         )
     )
     module = AnsibleModule(argument_spec=argument_spec)
+    return module
+
+
+def main():
+    module = create_snapshot_ansible_module()
 
     if not HAS_BOTO:
         module.fail_json(msg='boto required for this module')
@@ -155,91 +282,30 @@ def main():
     device_name = module.params.get('device_name')
     wait = module.params.get('wait')
     wait_timeout = module.params.get('wait_timeout')
-    snapshot_max_age = module.params.get('snapshot_max_age')
+    last_snapshot_min_age = module.params.get('last_snapshot_min_age')
     snapshot_tags = module.params.get('snapshot_tags')
     state = module.params.get('state')
 
-    snapshot = None
-    changed = False
-
-    if not volume_id and not instance_id and not snapshot_id or volume_id and instance_id and snapshot_id:
-        module.fail_json(msg='One and only one of volume_id or instance_id or snapshot_id must be specified')
-    if instance_id and not device_name or device_name and not instance_id:
-        module.fail_json(msg='Instance ID and device name must both be specified')
-
     ec2 = ec2_connect(module)
 
-    if instance_id:
+    create_snapshot(
-        try:
+        module=module,
-            volumes = ec2.get_all_volumes(filters={'attachment.instance-id': instance_id, 'attachment.device': device_name})
+        state=state,
-            if not volumes:
+        description=description,
-                module.fail_json(msg="Could not find volume with name %s attached to instance %s" % (device_name, instance_id))
+        wait=wait,
-            volume_id = volumes[0].id
+        wait_timeout=wait_timeout,
-        except boto.exception.BotoServerError, e:
+        ec2=ec2,
-            module.fail_json(msg = "%s: %s" % (e.error_code, e.error_message))
+        volume_id=volume_id,
-
+        instance_id=instance_id,
-    if state == 'absent':
+        snapshot_id=snapshot_id,
-        if not snapshot_id:
+        device_name=device_name,
-            module.fail_json(msg = 'snapshot_id must be set when state is absent')
+        snapshot_tags=snapshot_tags,
-        try:
+        last_snapshot_min_age=last_snapshot_min_age
-            snapshots = ec2.get_all_snapshots([snapshot_id])
+    )
-            ec2.delete_snapshot(snapshot_id)
-            module.exit_json(changed=True)
-        except boto.exception.BotoServerError, e:
-            # exception is raised if snapshot does not exist
-            if e.error_code == 'InvalidSnapshot.NotFound':
-                module.exit_json(changed=False)
-            else:
-                module.fail_json(msg = "%s: %s" % (e.error_code, e.error_message))
-
-    if snapshot_max_age > 0:
-        try:
-            snapshot_max_age = snapshot_max_age * 60 # Convert to seconds
-            current_snapshots = ec2.get_all_snapshots(filters={'volume_id': volume_id})
-            # Find the most recent snapshot
-            recent = dict(start_time=0, snapshot=None)
-            for s in current_snapshots:
-                start_time = time.mktime(time.strptime(s.start_time, '%Y-%m-%dT%H:%M:%S.000Z'))
-                if start_time > recent['start_time']:
-                    recent['start_time'] = start_time
-                    recent['snapshot'] = s
-
-            # Adjust snapshot start time to local timezone
-            tz_adjustment = time.daylight and time.altzone or time.timezone
-            recent['start_time'] -= tz_adjustment
-
-            # See if the snapshot is younger that the given max age
-            current_time = time.mktime(time.localtime())
-            snapshot_age = current_time - recent['start_time']
-            if snapshot_age < snapshot_max_age:
-                snapshot = recent['snapshot']
-        except boto.exception.BotoServerError, e:
-            module.fail_json(msg="%s: %s" % (e.error_code, e.error_message))
-
-    try:
-        # Create a new snapshot if we didn't find an existing one to use
-        if snapshot is None:
-            snapshot = ec2.create_snapshot(volume_id, description=description)
-            changed = True
-        if wait:
-            time_waited = 0
-            snapshot.update()
-            while snapshot.status != 'completed':
-                time.sleep(3)
-                snapshot.update()
-                time_waited += 3
-                if wait_timeout and time_waited > wait_timeout:
-                    module.fail_json('Timed out while creating snapshot.')
-        for k, v in snapshot_tags.items():
-            snapshot.add_tag(k, v)
-    except boto.exception.BotoServerError, e:
-        module.fail_json(msg="%s: %s" % (e.error_code, e.error_message))
-
-    module.exit_json(changed=changed, snapshot_id=snapshot.id, volume_id=snapshot.volume_id,
-            volume_size=snapshot.volume_size, tags=snapshot.tags.copy())
 
 # import module snippets
 from ansible.module_utils.basic import *
 from ansible.module_utils.ec2 import *
 
-main()
+if __name__ == '__main__':
+    main()