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,64 +128,87 @@ except ImportError:
     HAS_BOTO = False
 
 
-def main():
+# Find the most recent snapshot
-    argument_spec = ec2_argument_spec()
+def _get_snapshot_starttime(snap):
-    argument_spec.update(
+    return datetime.datetime.strptime(snap.start_time, '%Y-%m-%dT%H:%M:%S.000Z')
-        dict(
-            volume_id = dict(),
-            description = dict(),
-            instance_id = dict(),
-            snapshot_id = dict(),
-            device_name = dict(),
-            wait = dict(type='bool', default=True),
-            wait_timeout = dict(type='int', default=0),
-            snapshot_max_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)
 
-    if not HAS_BOTO:
-        module.fail_json(msg='boto required for this module')
 
-    volume_id = module.params.get('volume_id')
+def _get_most_recent_snapshot(snapshots, max_snapshot_age_secs=None, now=None):
-    snapshot_id = module.params.get('snapshot_id')
+    """
-    description = module.params.get('description')
+    Gets the most recently created snapshot and optionally filters the result
-    instance_id = module.params.get('instance_id')
+    if the snapshot is too old
-    device_name = module.params.get('device_name')
+    :param snapshots: list of snapshots to search
-    wait = module.params.get('wait')
+    :param max_snapshot_age_secs: filter the result if its older than this
-    wait_timeout = module.params.get('wait_timeout')
+    :param now: simulate time -- used for unit testing
-    snapshot_max_age = module.params.get('snapshot_max_age')
+    :return:
-    snapshot_tags = module.params.get('snapshot_tags')
+    """
-    state = module.params.get('state')
+    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
 
-    if not volume_id and not instance_id and not snapshot_id or volume_id and instance_id and snapshot_id:
+    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')
 
-    ec2 = ec2_connect(module)
-
     if instance_id:
         try:
             volumes = ec2.get_all_volumes(filters={'attachment.instance-id': instance_id, 'attachment.device': device_name})
-            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
         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:
-            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':
@@ -192,54 +216,96 @@ def main():
             else:
                 module.fail_json(msg = "%s: %s" % (e.error_code, e.error_message))
 
-    if snapshot_max_age > 0:
+        # successful delete
+        module.exit_json(changed=True)
+
+    if last_snapshot_min_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))
 
+        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:
-            time_waited = 0
+            if not _create_with_wait(snapshot, wait_timeout):
-            snapshot.update()
+                module.fail_json(msg='Timed out while creating snapshot.')
-            while snapshot.status != 'completed':
+        if snapshot_tags:
-                time.sleep(3)
+            for k, v in snapshot_tags.items():
-                snapshot.update()
+                snapshot.add_tag(k, v)
-                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,
+    module.exit_json(changed=changed,
-            volume_size=snapshot.volume_size, tags=snapshot.tags.copy())
+                     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(
+            volume_id = dict(),
+            description = dict(),
+            instance_id = dict(),
+            snapshot_id = dict(),
+            device_name = dict(),
+            wait = dict(type='bool', default=True),
+            wait_timeout = 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')
+
+    volume_id = module.params.get('volume_id')
+    snapshot_id = module.params.get('snapshot_id')
+    description = module.params.get('description')
+    instance_id = module.params.get('instance_id')
+    device_name = module.params.get('device_name')
+    wait = module.params.get('wait')
+    wait_timeout = module.params.get('wait_timeout')
+    last_snapshot_min_age = module.params.get('last_snapshot_min_age')
+    snapshot_tags = module.params.get('snapshot_tags')
+    state = module.params.get('state')
+
+    ec2 = ec2_connect(module)
+
+    create_snapshot(
+        module=module,
+        state=state,
+        description=description,
+        wait=wait,
+        wait_timeout=wait_timeout,
+        ec2=ec2,
+        volume_id=volume_id,
+        instance_id=instance_id,
+        snapshot_id=snapshot_id,
+        device_name=device_name,
+        snapshot_tags=snapshot_tags,
+        last_snapshot_min_age=last_snapshot_min_age
+    )
 
 # import module snippets
 from ansible.module_utils.basic import *
 from ansible.module_utils.ec2 import *
 
-main()
+if __name__ == '__main__':
+    main()