Fix tests, new tests, code cleanup

Author: heerener

hpc_provisioner/src/hpc_provisioner/aws_queries.py

@@ -6,6 +6,12 @@
 import boto3
 from botocore.exceptions import ClientError
 
+from hpc_provisioner.constants import (
+    BILLING_TAG_KEY,
+    BILLING_TAG_VALUE,
+    PROJECT_TAG_KEY,
+    VLAB_TAG_KEY,
+)
 from hpc_provisioner.dynamodb_actions import (
     SubnetAlreadyRegisteredException,
     dynamodb_client,
@@ -61,6 +67,31 @@ def create_keypair(ec2_client, vlab_id, project_id, tags) -> dict:
         )
 
 
+def create_secret(sm_client, vlab_id, project_id, secret_name, secret_value):
+    secret = sm_client.create_secret(
+        Name=secret_name,
+        Description=f"SSH Key for cluster for vlab {vlab_id}, project {project_id}",
+        SecretString=secret_value,
+        Tags=[
+            {"Key": VLAB_TAG_KEY, "Value": vlab_id},
+            {"Key": PROJECT_TAG_KEY, "Value": project_id},
+            {"Key": BILLING_TAG_KEY, "Value": BILLING_TAG_VALUE},
+        ],
+    )
+
+    return secret
+
+
+def get_secret(sm_client, secret_name):
+    existing_secrets = sm_client.list_secrets(Filters=[{"Key": "name", "Values": [secret_name]}])
+    if secret_list := existing_secrets.get("SecretList", []):
+        secret = secret_list[0]
+    else:
+        raise RuntimeError(f"Secret {secret_name} does not exist in SecretsManager")
+
+    return secret
+
+
 def get_efs(efs_client) -> str:
     """
     Get the ID for the EFS for pclusters

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
+from hpc_provisioner.aws_queries import create_keypair, create_secret, get_secret
 from hpc_provisioner.constants import (
     BILLING_TAG_KEY,
     BILLING_TAG_VALUE,
@@ -80,27 +80,11 @@ def pcluster_create_request_handler(event, _context=None):
     )
 
     if "KeyMaterial" in ssh_keypair:
-        secret = sm_client.create_secret(
-            Name=ssh_keypair["KeyName"],
-            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},
-            ],
+        secret = create_secret(
+            sm_client, vlab_id, project_id, ssh_keypair["KeyName"], ssh_keypair["KeyMaterial"]
         )
     else:
-        existing_secrets = sm_client.list_secrets(
-            Filters=[{"Key": "name", "Values": [ssh_keypair["KeyName"]]}]
-        )
-        if secret_list := existing_secrets["SecretList"]:
-            secret = secret_list[0]
-        else:
-            raise RuntimeError(
-                f"SSH Keypair {ssh_keypair['KeyName']} already exists in EC2 but "
-                "was not stored in SecretsManager - unable to retrieve private key"
-            )
+        secret = get_secret(sm_client, ssh_keypair["KeyName"])
 
     logger.debug("calling create lambda async")
     boto3.client("lambda").invoke_async(

hpc_provisioner/tests/test_aws_queries.py

@@ -2,17 +2,25 @@
 from unittest.mock import MagicMock, call, patch
 
 import pytest
+from botocore.exceptions import ClientError
 from hpc_provisioner.aws_queries import (
     CouldNotDetermineEFSException,
-    CouldNotDetermineKeyPairException,
     CouldNotDetermineSecurityGroupException,
     OutOfSubnetsException,
     claim_subnet,
+    create_keypair,
+    create_secret,
     get_available_subnet,
     get_efs,
-    get_keypair,
+    get_secret,
     get_security_group,
 )
+from hpc_provisioner.constants import (
+    BILLING_TAG_KEY,
+    BILLING_TAG_VALUE,
+    PROJECT_TAG_KEY,
+    VLAB_TAG_KEY,
+)
 from hpc_provisioner.dynamodb_actions import SubnetAlreadyRegisteredException
 
 logger = logging.getLogger("test_logger")
@@ -28,36 +36,6 @@
 logger.setLevel(logging.DEBUG)
 
 
-@pytest.mark.parametrize(
-    "keypairs",
-    [
-        {"KeyPairs": [{"KeyName": "keypair-1"}]},
-    ],
-)
-def test_get_keypair(keypairs):
-    mock_ec2_client = MagicMock()
-    mock_ec2_client.describe_key_pairs.return_value = keypairs
-    keypair = get_keypair(mock_ec2_client)
-    assert keypair == keypairs["KeyPairs"][0]["KeyName"]
-
-
-@pytest.mark.parametrize(
-    "keypairs",
-    [
-        {"KeyPairs": []},
-        {"KeyPairs": ["keypair-1", "keypair-2"]},
-    ],
-)
-def test_get_keypair_fails(keypairs):
-    mock_ec2_client = MagicMock()
-    mock_ec2_client.describe_key_pairs.return_value = keypairs
-    with pytest.raises(
-        CouldNotDetermineKeyPairException,
-        match=str(keypairs["KeyPairs"]).replace("[", "\\["),
-    ):
-        get_keypair(mock_ec2_client)
-
-
 @pytest.mark.parametrize(
     "filesystems",
     [
@@ -333,3 +311,54 @@ def test_get_available_subnet(mock_dynamodb_client, mock_claim_subnet):
         mock_dynamodb_client(), ec2_subnets["Subnets"], cluster_name
     )
     assert subnet == "sub-1"
+
+
+def test_create_keypair():
+    mock_ec2_client = MagicMock()
+    mock_ec2_client.describe_key_pairs.side_effect = ClientError(
+        error_response={"Error": {"Code": 1, "Message": "It failed"}},
+        operation_name="describe_key_pairs",
+    )
+    mock_ec2_client.create_key_pair.return_value = "key created"
+    vlab_id = "test_vlab"
+    project_id = "test_project"
+    tags = [{"Key": "tagkey", "Value": "tagvalue"}]
+    create_keypair(mock_ec2_client, vlab_id, project_id, tags)
+    mock_ec2_client.create_key_pair.assert_called_once_with(
+        KeyName=f"pcluster-{vlab_id}-{project_id}",
+        TagSpecifications=[{"ResourceType": "key-pair", "Tags": tags}],
+    )
+
+
+def test_get_secret():
+    secret_value = "supersecret"
+    mock_sm_client = MagicMock()
+    mock_sm_client.list_secrets.return_value = {"SecretList": [secret_value]}
+    retrieved_secret = get_secret(mock_sm_client, "mysecret")
+    assert retrieved_secret == secret_value
+
+
+def test_get_secret_not_found():
+    mock_sm_client = MagicMock()
+    mock_sm_client.list_secrets.return_value = {}
+    with pytest.raises(RuntimeError):
+        get_secret(mock_sm_client, "mysecret")
+
+
+def test_create_secret():
+    mock_sm_client = MagicMock()
+    vlab_id = "test_vlab"
+    project_id = "test_project"
+    secret_name = "mysecret"
+    secret_value = "supersecret"
+    create_secret(mock_sm_client, vlab_id, project_id, secret_name, secret_value)
+    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=secret_value,
+        Tags=[
+            {"Key": VLAB_TAG_KEY, "Value": vlab_id},
+            {"Key": PROJECT_TAG_KEY, "Value": project_id},
+            {"Key": BILLING_TAG_KEY, "Value": BILLING_TAG_VALUE},
+        ],
+    )

hpc_provisioner/tests/test_resource_provisioner.py

@@ -156,22 +156,48 @@ def test_get_all_clusters(data):
 
 
 @patch("hpc_provisioner.handlers.boto3")
-def test_post(patched_boto3, post_event):
+@pytest.mark.parametrize("key_exists", [True, False])
+def test_post(patched_boto3, post_event, key_exists):
+    test_cluster_name = cluster_name(post_event["vlab_id"], post_event["project_id"])
     mock_client = MagicMock()
     patched_boto3.client.return_value = mock_client
-    actual_response = handlers.pcluster_create_request_handler(post_event)
-    mock_client.invoke_async.assert_called_once_with(
+    with patch("hpc_provisioner.handlers.create_keypair") as patched_create_keypair:
+        if key_exists:
+            patched_create_keypair.return_value = {
+                "KeyName": test_cluster_name,
+            }
+        else:
+            patched_create_keypair.return_value = {
+                "KeyMaterial": "secret_stuff",
+                "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)
+    mock_client.invoke_async.assert_called_with(
         FunctionName="hpc-resource-provisioner-creator",
         InvokeArgs=json.dumps(
-            {"vlab_id": post_event["vlab_id"], "project_id": post_event["project_id"]}
+            {
+                "vlab_id": post_event["vlab_id"],
+                "project_id": post_event["project_id"],
+                "keyname": f"pcluster-{post_event['vlab_id']}-{post_event['project_id']}",
+            }
         ),
     )
     expected_response = expected_response_template(
         text=json.dumps(
             {
                 "cluster": {
-                    "clusterName": cluster_name(post_event["vlab_id"], post_event["project_id"]),
+                    "clusterName": test_cluster_name,
                     "clusterStatus": "CREATE_REQUEST_RECEIVED",
+                    "private_ssh_key_arn": "secret ARN",
                 }
             }
         )
@@ -183,7 +209,10 @@ def test_post(patched_boto3, post_event):
     "hpc_provisioner.aws_queries.dynamodb_client",
 )
 @patch("hpc_provisioner.aws_queries.free_subnet")
-def test_delete(patched_free_subnet, patched_dynamodb_client, data, delete_event):
+@patch("hpc_provisioner.pcluster_manager.remove_key")
+def test_delete(
+    patched_remove_key, patched_free_subnet, patched_dynamodb_client, data, delete_event
+):
     mock_client = MagicMock()
     patched_dynamodb_client.return_value = mock_client
     with patch(
@@ -204,6 +233,7 @@ def test_delete(patched_free_subnet, patched_dynamodb_client, data, delete_event
     expected_response = expected_response_template(text=json.dumps(data["deletingCluster"]))
     assert actual_response == expected_response
     patched_get_registered_subnets.assert_called_once()
+    patched_remove_key.assert_called_once()
     call1 = call(mock_client, "subnet-123")
     call2 = call(mock_client, "subnet-234")
     patched_free_subnet.assert_has_calls([call1, call2], any_order=True)
@@ -235,7 +265,8 @@ def test_get_internal_server_error(get_event):
 @patch(
     "hpc_provisioner.aws_queries.dynamodb_client",
 )
-def test_delete_not_found(patched_dynamodb_client, delete_event):
+@patch("hpc_provisioner.pcluster_manager.remove_key")
+def test_delete_not_found(patched_remove_key, patched_dynamodb_client, delete_event):
     error_message = f"Cluster {delete_event['vlab_id']}-{delete_event['project_id']} does not exist"
     with patch(
         "hpc_provisioner.pcluster_manager.pc.delete_cluster",
@@ -245,12 +276,14 @@ def test_delete_not_found(patched_dynamodb_client, delete_event):
         delete_cluster.assert_called_once()
     assert result == {"statusCode": 404, "body": error_message}
     patched_dynamodb_client.assert_called_once()
+    patched_remove_key.assert_called_once()
 
 
 @patch(
     "hpc_provisioner.aws_queries.dynamodb_client",
 )
-def test_delete_internal_server_error(patched_dynamodb_client, delete_event):
+@patch("hpc_provisioner.pcluster_manager.remove_key")
+def test_delete_internal_server_error(patched_remove_key, patched_dynamodb_client, delete_event):
     with patch(
         "hpc_provisioner.pcluster_manager.pc.delete_cluster",
         side_effect=RuntimeError,
@@ -259,6 +292,7 @@ def test_delete_internal_server_error(patched_dynamodb_client, delete_event):
         patched_delete_cluster.assert_called_once()
     assert result == {"statusCode": 500, "body": "<class 'RuntimeError'>"}
     patched_dynamodb_client.assert_called_once()
+    patched_remove_key.assert_called_once()
 
 
 @patch("hpc_provisioner.pcluster_manager.pc.create_cluster")
@@ -278,9 +312,7 @@ def test_do_create_already_exists(patched_boto3, patched_create_cluster, post_ev
 @patch("hpc_provisioner.pcluster_manager.get_available_subnet", return_value="subnet-123")
 @patch("hpc_provisioner.pcluster_manager.get_security_group", return_value="sg-123")
 @patch("hpc_provisioner.pcluster_manager.get_efs", return_value="efs-123")
-@patch("hpc_provisioner.pcluster_manager.get_keypair", return_value="keypair-123")
 def test_do_create(
-    patched_get_keypair,
     patched_get_efs,
     patched_get_security_group,
     patched_get_available_subnet,
@@ -299,13 +331,13 @@ def test_do_create(
         "ec2": mock_ec2_client,
         "efs": mock_efs_client,
     }[x]
+    post_event["keyname"] = cluster_name(post_event["vlab_id"], post_event["project_id"])
     handlers.pcluster_do_create_handler(post_event)
     patched_create_cluster.assert_called_once()
     assert patched_create_cluster.call_args.kwargs["cluster_name"] == cluster_name(
         post_event["vlab_id"], post_event["project_id"]
     )
     assert "tmp" in patched_create_cluster.call_args.kwargs["cluster_configuration"]
-    patched_get_keypair.assert_called_once()
     patched_get_efs.assert_called_once()
     patched_get_security_group.assert_called_once()
     patched_get_available_subnet.assert_called_once()