Refactor ec2_key* modules (#2168)

SUMMARY
Refactor ec2_key, ec2_key_info modules
ISSUE TYPE


Feature Pull Request

COMPONENT NAME


ec2_key
ec2_key_info

Reviewed-by: Alina Buzachis

changelogs/fragments/20240107-refactor_ec2_key-modules.yml

@@ -0,0 +1,4 @@
+---
+minor_changes:
+  - ec2_key - refactored code to use ``AnsibleEC2Error`` as well as moving shared code into module_utils.ec2 (https://github.com/ansible-collections/amazon.aws/pull/2168).
+  - ec2_key_info - refactored code to use ``AnsibleEC2Error`` as well as moving shared code into module_utils.ec2 (https://github.com/ansible-collections/amazon.aws/pull/2168).

plugins/modules/ec2_key.py

@@ -164,17 +164,15 @@
 import os
 import uuid
 
-try:
-    import botocore
-except ImportError:
-    pass  # caught by AnsibleAWSModule
-
 from ansible.module_utils._text import to_bytes
 
-from ansible_collections.amazon.aws.plugins.module_utils.botocore import is_boto3_error_code
+from ansible_collections.amazon.aws.plugins.module_utils.ec2 import AnsibleEC2Error
+from ansible_collections.amazon.aws.plugins.module_utils.ec2 import create_key_pair
+from ansible_collections.amazon.aws.plugins.module_utils.ec2 import delete_key_pair as delete_ec2_key_pair
+from ansible_collections.amazon.aws.plugins.module_utils.ec2 import describe_key_pairs
 from ansible_collections.amazon.aws.plugins.module_utils.ec2 import ensure_ec2_tags
+from ansible_collections.amazon.aws.plugins.module_utils.ec2 import import_key_pair
 from ansible_collections.amazon.aws.plugins.module_utils.modules import AnsibleAWSModule
-from ansible_collections.amazon.aws.plugins.module_utils.retries import AWSRetry
 from ansible_collections.amazon.aws.plugins.module_utils.tagging import boto3_tag_list_to_ansible_dict
 from ansible_collections.amazon.aws.plugins.module_utils.tagging import boto3_tag_specifications
 from ansible_collections.amazon.aws.plugins.module_utils.transformation import scrub_none_parameters
@@ -193,9 +191,9 @@ def _import_key_pair(ec2_client, name, key_material, tag_spec=None):
     params = scrub_none_parameters(params)
 
     try:
-        key = ec2_client.import_key_pair(aws_retry=True, **params)
-    except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as err:
-        raise Ec2KeyFailure(err, "error importing key")
+        key = import_key_pair(ec2_client, **params)
+    except AnsibleEC2Error as e:
+        raise Ec2KeyFailure(e, "error importing key")
     return key
 
 
@@ -236,18 +234,15 @@ def get_key_fingerprint(check_mode, ec2_client, key_material):
 
 def find_key_pair(ec2_client, name):
     try:
-        key = ec2_client.describe_key_pairs(aws_retry=True, KeyNames=[name])
-    except is_boto3_error_code("InvalidKeyPair.NotFound"):
-        return None
-    except (
-        botocore.exceptions.ClientError,
-        botocore.exceptions.BotoCoreError,
-    ) as err:  # pylint: disable=duplicate-except
-        raise Ec2KeyFailure(err, "error finding keypair")
+        key = describe_key_pairs(ec2_client, KeyNames=[name])
+        if not key:
+            return None
+    except AnsibleEC2Error as e:
+        raise Ec2KeyFailure(e, "error finding keypair")
     except IndexError:
         key = None
 
-    return key["KeyPairs"][0]
+    return key[0]
 
 
 def _create_key_pair(ec2_client, name, tag_spec, key_type):
@@ -260,9 +255,9 @@ def _create_key_pair(ec2_client, name, tag_spec, key_type):
     params = scrub_none_parameters(params)
 
     try:
-        key = ec2_client.create_key_pair(aws_retry=True, **params)
-    except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as err:
-        raise Ec2KeyFailure(err, "error creating key")
+        key = create_key_pair(ec2_client, **params)
+    except AnsibleEC2Error as e:
+        raise Ec2KeyFailure(e, "error creating key")
     return key
 
 
@@ -325,13 +320,6 @@ def update_key_pair_by_key_type(check_mode, ec2_client, name, key_type, tag_spec
         return {"changed": True, "key": key_data, "msg": "key pair updated"}
 
 
-def _delete_key_pair(ec2_client, key_name):
-    try:
-        ec2_client.delete_key_pair(aws_retry=True, KeyName=key_name)
-    except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as err:
-        raise Ec2KeyFailure(err, "error deleting key")
-
-
 def delete_key_pair(check_mode, ec2_client, name, finish_task=True):
     key = find_key_pair(ec2_client, name)
 
@@ -341,7 +329,10 @@ def delete_key_pair(check_mode, ec2_client, name, finish_task=True):
         result = {"key": None, "msg": "key did not exist"}
         return result
     else:
-        _delete_key_pair(ec2_client, name)
+        try:
+            delete_ec2_key_pair(ec2_client, name)
+        except AnsibleEC2Error as e:
+            raise Ec2KeyFailure(e, "error deleting keypair")
         if not finish_task:
             return
         result = {"changed": True, "key": None, "msg": "key deleted"}
@@ -389,7 +380,7 @@ def main():
         supports_check_mode=True,
     )
 
-    ec2_client = module.client("ec2", retry_decorator=AWSRetry.jittered_backoff())
+    ec2_client = module.client("ec2")
 
     name = module.params["name"]
     state = module.params.get("state")

plugins/modules/ec2_key_info.py

@@ -114,9 +114,9 @@
 
 from ansible.module_utils.common.dict_transformations import camel_dict_to_snake_dict
 
-from ansible_collections.amazon.aws.plugins.module_utils.botocore import is_boto3_error_code
+from ansible_collections.amazon.aws.plugins.module_utils.ec2 import AnsibleEC2Error
+from ansible_collections.amazon.aws.plugins.module_utils.ec2 import describe_key_pairs
 from ansible_collections.amazon.aws.plugins.module_utils.modules import AnsibleAWSModule
-from ansible_collections.amazon.aws.plugins.module_utils.retries import AWSRetry
 from ansible_collections.amazon.aws.plugins.module_utils.tagging import boto3_tag_list_to_ansible_dict
 from ansible_collections.amazon.aws.plugins.module_utils.transformation import ansible_dict_to_boto3_filter_list
 
@@ -140,14 +140,12 @@ def list_ec2_key_pairs(connection, module):
         params["IncludePublicKey"] = True
 
     try:
-        result = connection.describe_key_pairs(**params)
-    except is_boto3_error_code("InvalidKeyPair.NotFound"):
-        result = {}
-    except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
+        result = describe_key_pairs(connection, **params)
+    except AnsibleEC2Error as e:
         module.fail_json_aws(e, msg="Failed to list EC2 key pairs")
 
     # Turn the boto3 result in to ansible_friendly_snaked_names
-    snaked_keys = [camel_dict_to_snake_dict(key) for key in result.get("KeyPairs", [])]
+    snaked_keys = [camel_dict_to_snake_dict(key) for key in result]
 
     # Turn the boto3 result in to ansible friendly tag dictionary
     for key in snaked_keys:
@@ -170,7 +168,7 @@ def main():
     )
 
     try:
-        connection = module.client("ec2", retry_decorator=AWSRetry.jittered_backoff())
+        connection = module.client("ec2")
     except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
         module.fail_json_aws(e, msg="Failed to connect to AWS")
 

tests/unit/plugins/modules/test_ec2_key.py

@@ -3,7 +3,6 @@
 
 import copy
 import datetime
-from unittest.mock import ANY
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -40,7 +39,8 @@ def raise_botocore_exception_clienterror(action):
     return botocore.exceptions.ClientError(params, action)
 
 
-def test__import_key_pair():
+@patch(module_name + ".import_key_pair")
+def test__import_key_pair(m_import_key_pair):
     ec2_client = MagicMock()
     name = "my_keypair"
     key_material = "ssh-rsa AAAAB3NzaC1yc2EAA [email protected]"
@@ -50,17 +50,17 @@ def test__import_key_pair():
         "PublicKeyMaterial": to_bytes(key_material),
     }
 
-    ec2_client.import_key_pair.return_value = {
+    m_import_key_pair.return_value = {
         "KeyFingerprint": "d7:ff:a6:63:18:64:9c:57:a1:ee:ca:a4:ad:c2:81:62",
         "KeyName": "my_keypair",
         "KeyPairId": "key-012345678905a208d",
     }
 
     result = ec2_key._import_key_pair(ec2_client, name, key_material)
 
-    assert result == ec2_client.import_key_pair.return_value
-    assert ec2_client.import_key_pair.call_count == 1
-    ec2_client.import_key_pair.assert_called_with(aws_retry=True, **expected_params)
+    assert result == m_import_key_pair.return_value
+    assert m_import_key_pair.call_count == 1
+    m_import_key_pair.assert_called_with(ec2_client, **expected_params)
 
 
 def test_api_failure__import_key_pair():
@@ -155,27 +155,26 @@ def test_get_key_fingerprint(m_find_key_pair, m_import_key_pair, m_delete_key_pa
     assert m_delete_key_pair.call_count == 1
 
 
-def test_find_key_pair():
+@patch(module_name + ".describe_key_pairs")
+def test_find_key_pair(m_describe_key_pairs):
     ec2_client = MagicMock()
     name = "my_keypair"
 
-    ec2_client.describe_key_pairs.return_value = {
-        "KeyPairs": [
-            {
-                "CreateTime": datetime.datetime(2022, 9, 15, 20, 10, 15, tzinfo=tzutc()),
-                "KeyFingerprint": "11:12:13:14:bb:26:85:b2:e8:39:27:bc:ee:aa:ff:ee:dd:cc:bb:aa",
-                "KeyName": "my_keypair",
-                "KeyPairId": "key-043046ef2a9a80b56",
-                "KeyType": "rsa",
-                "Tags": [],
-            }
-        ],
-    }
+    m_describe_key_pairs.return_value = [
+        {
+            "CreateTime": datetime.datetime(2022, 9, 15, 20, 10, 15, tzinfo=tzutc()),
+            "KeyFingerprint": "11:12:13:14:bb:26:85:b2:e8:39:27:bc:ee:aa:ff:ee:dd:cc:bb:aa",
+            "KeyName": "my_keypair",
+            "KeyPairId": "key-043046ef2a9a80b56",
+            "KeyType": "rsa",
+            "Tags": [],
+        }
+    ]
 
     ec2_key.find_key_pair(ec2_client, name)
 
-    assert ec2_client.describe_key_pairs.call_count == 1
-    ec2_client.describe_key_pairs.assert_called_with(aws_retry=True, KeyNames=[name])
+    assert m_describe_key_pairs.call_count == 1
+    m_describe_key_pairs.assert_called_with(ec2_client, KeyNames=[name])
 
 
 def test_api_failure_find_key_pair():
@@ -199,15 +198,16 @@ def test_invalid_key_pair_find_key_pair():
     assert result is None
 
 
-def test__create_key_pair():
+@patch(module_name + ".create_key_pair")
+def test__create_key_pair(m_create_key_pair):
     ec2_client = MagicMock()
     name = "my_keypair"
     tag_spec = None
     key_type = None
 
     expected_params = {"KeyName": name}
 
-    ec2_client.create_key_pair.return_value = {
+    m_create_key_pair.return_value = {
         "KeyFingerprint": "d7:ff:a6:63:18:64:9c:57:a1:ee:ca:a4:ad:c2:81:62",
         "KeyMaterial": (
             "-----BEGIN RSA PRIVATE KEY-----\n"  # gitleaks:allow
@@ -220,9 +220,9 @@ def test__create_key_pair():
 
     result = ec2_key._create_key_pair(ec2_client, name, tag_spec, key_type)
 
-    assert result == ec2_client.create_key_pair.return_value
-    assert ec2_client.create_key_pair.call_count == 1
-    ec2_client.create_key_pair.assert_called_with(aws_retry=True, **expected_params)
+    assert result == m_create_key_pair.return_value
+    assert m_create_key_pair.call_count == 1
+    m_create_key_pair.assert_called_with(ec2_client, **expected_params)
 
 
 def test_api_failure__create_key_pair():
@@ -310,26 +310,6 @@ def test_create_new_key_pair_no_key_material(m_create_key_pair, m_extract_key_da
     assert m_extract_key_data.call_count == 1
 
 
-def test__delete_key_pair():
-    ec2_client = MagicMock()
-
-    key_name = "my_keypair"
-    ec2_key._delete_key_pair(ec2_client, key_name)
-
-    assert ec2_client.delete_key_pair.call_count == 1
-    ec2_client.delete_key_pair.assert_called_with(aws_retry=True, KeyName=key_name)
-
-
-def test_api_failure__delete_key_pair():
-    ec2_client = MagicMock()
-    name = "my_keypair"
-
-    ec2_client.delete_key_pair.side_effect = raise_botocore_exception_clienterror("delete_key_pair")
-
-    with pytest.raises(ec2_key.Ec2KeyFailure):
-        ec2_key._delete_key_pair(ec2_client, name)
-
-
 @patch(module_name + ".extract_key_data")
 @patch(module_name + "._import_key_pair")
 @patch(module_name + ".delete_key_pair")
@@ -577,7 +557,7 @@ def test_handle_existing_key_pair_else(m_extract_key_data):
     assert m_extract_key_data.call_count == 1
 
 
-@patch(module_name + "._delete_key_pair")
+@patch(module_name + ".delete_ec2_key_pair")
 @patch(module_name + ".find_key_pair")
 def test_delete_key_pair_key_exists(m_find_key_pair, m_delete_key_pair, tmp_path):
     module = MagicMock()
@@ -609,7 +589,7 @@ def test_delete_key_pair_key_exists(m_find_key_pair, m_delete_key_pair, tmp_path
     assert result == {"changed": True, "key": None, "msg": "key deleted"}
 
 
-@patch(module_name + "._delete_key_pair")
+@patch(module_name + ".delete_ec2_key_pair")
 @patch(module_name + ".find_key_pair")
 def test_delete_key_pair_key_not_exist(m_find_key_pair, m_delete_key_pair):
     module = MagicMock()
@@ -656,4 +636,4 @@ def test_main_success(m_AnsibleAWSModule):
 
     ec2_key.main()
 
-    m_module.client.assert_called_with("ec2", retry_decorator=ANY)
+    m_module.client.assert_called_with("ec2")