Refactor: remove some logic from handlers into helpers

Author: heerener

hpc_provisioner/src/hpc_provisioner/aws_queries.py

@@ -42,10 +42,6 @@ class CouldNotDetermineEFSException(Exception):
     """
 
 
-class CouldNotDetermineKeyPairException(Exception):
-    """Indicates that we either found too many or no keypairs with tags HPC_Goal:compute_cluster"""
-
-
 def get_cluster_name(vlab_id: str, project_id: str) -> str:
     return f"pcluster-{vlab_id}-{project_id}"
 
@@ -67,6 +63,17 @@ def create_keypair(ec2_client, vlab_id, project_id, tags) -> dict:
         )
 
 
+def store_private_key(sm_client, vlab_id, project_id, ssh_keypair):
+    if "KeyMaterial" in ssh_keypair:
+        secret = create_secret(
+            sm_client, vlab_id, project_id, ssh_keypair["KeyName"], ssh_keypair["KeyMaterial"]
+        )
+    else:
+        secret = get_secret(sm_client, ssh_keypair["KeyName"])
+
+    return secret
+
+
 def create_secret(sm_client, vlab_id, project_id, secret_name, secret_value):
     secret = sm_client.create_secret(
         Name=secret_name,

hpc_provisioner/src/hpc_provisioner/handlers.py

@@ -6,7 +6,7 @@
 import boto3
 from pcluster.api.errors import NotFoundException
 
-from hpc_provisioner.aws_queries import create_keypair, create_secret, get_secret
+from hpc_provisioner.aws_queries import create_keypair, store_private_key
 from hpc_provisioner.constants import (
     BILLING_TAG_KEY,
     BILLING_TAG_VALUE,
@@ -79,12 +79,7 @@ def pcluster_create_request_handler(event, _context=None):
         ],
     )
 
-    if "KeyMaterial" in ssh_keypair:
-        secret = create_secret(
-            sm_client, vlab_id, project_id, ssh_keypair["KeyName"], ssh_keypair["KeyMaterial"]
-        )
-    else:
-        secret = get_secret(sm_client, ssh_keypair["KeyName"])
+    secret = store_private_key(sm_client, vlab_id, project_id, ssh_keypair)
 
     logger.debug("calling create lambda async")
     boto3.client("lambda").invoke_async(

hpc_provisioner/tests/test_aws_queries.py

@@ -14,6 +14,7 @@
     get_efs,
     get_secret,
     get_security_group,
+    store_private_key,
 )
 from hpc_provisioner.constants import (
     BILLING_TAG_KEY,
@@ -362,3 +363,46 @@ def test_create_secret():
             {"Key": BILLING_TAG_KEY, "Value": BILLING_TAG_VALUE},
         ],
     )
+
+
+@pytest.mark.parametrize("keypair_exists", [True, False])
+@pytest.mark.parametrize("secret_exists", [True, False])
+def test_store_private_key(keypair_exists, secret_exists):
+    if secret_exists and not keypair_exists:
+        pytest.skip(
+            "New keypair with existing secret: sm_client.create_secret will raise "
+            "on trying to create a secret that already exists."
+        )
+    mock_sm_client = MagicMock()
+    vlab_id = "testvlab"
+    project_id = "testproject"
+    secret_name = key_name = f"pcluster-{vlab_id}-{project_id}"
+    if not keypair_exists:
+        ssh_keypair = {"KeyMaterial": "supersecret", "KeyName": key_name}
+    else:
+        ssh_keypair = {"KeyName": key_name}
+
+    if not keypair_exists:
+        if not secret_exists:
+            print("Keypair created, secret does not exist yet")
+            store_private_key(mock_sm_client, vlab_id, project_id, ssh_keypair)
+            mock_sm_client.create_secret.assert_called_once_with(
+                Name=secret_name,
+                Description=f"SSH Key for cluster for vlab {vlab_id}, project {project_id}",
+                SecretString=ssh_keypair["KeyMaterial"],
+                Tags=[
+                    {"Key": VLAB_TAG_KEY, "Value": vlab_id},
+                    {"Key": PROJECT_TAG_KEY, "Value": project_id},
+                    {"Key": BILLING_TAG_KEY, "Value": BILLING_TAG_VALUE},
+                ],
+            )
+    elif secret_exists:
+        print("Both already exist")
+        mock_sm_client.list_secrets.return_value = {"SecretList": ["somesecret"]}
+        retrieved_secret = store_private_key(mock_sm_client, vlab_id, project_id, ssh_keypair)
+        assert retrieved_secret == "somesecret"
+    else:
+        print("Keypair already existed but was not stored in secretsmanager yet")
+        mock_sm_client.list_secrets.return_value = {"SecretList": []}
+        with pytest.raises(RuntimeError):
+            store_private_key(mock_sm_client, vlab_id, project_id, ssh_keypair)

hpc_provisioner/tests/test_resource_provisioner.py

@@ -172,15 +172,9 @@ def test_post(patched_boto3, post_event, key_exists):
                 "KeyName": test_cluster_name,
             }
 
-        with patch("hpc_provisioner.handlers.create_secret") as patched_create_secret:
-            patched_create_keypair.return_value = {
-                "KeyMaterial": "secret_stuff",
-                "KeyName": test_cluster_name,
-            }
-            patched_create_secret.return_value = {"ARN": "secret ARN"}
-            with patch("hpc_provisioner.handlers.get_secret") as patched_get_secret:
-                patched_get_secret.return_value = {"ARN": "secret ARN"}
-                actual_response = handlers.pcluster_create_request_handler(post_event)
+        with patch("hpc_provisioner.handlers.store_private_key") as patched_store_private_key:
+            patched_store_private_key.return_value = {"ARN": "secret ARN"}
+            actual_response = handlers.pcluster_create_request_handler(post_event)
     mock_client.invoke_async.assert_called_with(
         FunctionName="hpc-resource-provisioner-creator",
         InvokeArgs=json.dumps(