Skip to content

Commit adc3aff

Browse files
author
Himani Anil Deshpande
committed
Add unit test for share_compute_fleet_dna.py
1 parent ea3c566 commit adc3aff

File tree

4 files changed

+210
-0
lines changed

4 files changed

+210
-0
lines changed

‎pytest.ini‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
pythonpath =
1414
cookbooks/aws-parallelcluster-platform/files/dcv
1515
cookbooks/aws-parallelcluster-environment/files/cloudwatch
16+
cookbooks/aws-parallelcluster-environment/files/cfn_hup_configuration
1617
cookbooks/aws-parallelcluster-computefleet/files/compute_fleet_status
1718
cookbooks/aws-parallelcluster-computefleet/files/clusterstatusmgtd
1819
cookbooks/aws-parallelcluster-environment/files/custom_action_executor

‎test/unit/cfn_hup_configuration/test_share_compute_fleet_dna.py‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
4+
# with the License. A copy of the License is located at
5+
#
6+
# http://aws.amazon.com/apache2.0/
7+
#
8+
# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
9+
# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
10+
# limitations under the License.
11+
import os
12+
from unittest.mock import MagicMock, patch
13+
import boto3
14+
import pytest
15+
import json
16+
from assertpy import assert_that
17+
from botocore.stub import Stubber
18+
from share_compute_fleet_dna import get_compute_launch_template_ids, get_user_data, get_write_directives_section, parse_proxy_config
19+
from base64 import b64encode
20+
21+
22+
@pytest.mark.parametrize(
23+
("launch_template_config_content", "errors"),
24+
[
25+
("{\"Queues\":{\"queue-0\":{\"ComputeResources\":{\"compute-resource-1\":{\"LaunchTemplate\":{\"Version\":\"1\",\"Id\":\"lt-037123456747c3bc5\"}},\"compute-resource-2\":{\"LaunchTemplate\":{\"Version\":\"1\",\"Id\":\"lt-0fcecb59a3721c0b3\"}},\"compute-resource-0\":{\"LaunchTemplate\":{\"Version\":\"1\",\"Id\":\"lt-12345678901234567\"}}}}}}", False),
26+
("{\"Queues\":{\"queue-0\":}}}", True),
27+
28+
],
29+
)
30+
def test_get_compute_launch_template_ids(mocker, launch_template_config_content, errors):
31+
mocker.patch("builtins.open", mocker.mock_open(read_data=launch_template_config_content))
32+
actual_op = get_compute_launch_template_ids(launch_template_config_content)
33+
if errors:
34+
assert_that(actual_op).is_none()
35+
else:
36+
assert_that(actual_op).is_equal_to(json.loads(launch_template_config_content))
37+
38+
39+
@pytest.mark.parametrize(
40+
("mime_user_data_file", "write_section"),
41+
[
42+
("user_data_1.txt",
43+
[{'path': '/tmp/dna.json', 'permissions': '0644', 'owner': 'root:root', 'content': '{"cluster":{"base_os":"alinux2","cluster_config_s3_key":"parallelcluster/clusters/dummy-cluster-randomstring123/configs/cluster-config-with-implied-values.yaml","cluster_config_version":"","cluster_name":"clustername","cluster_s3_bucket":"parallelcluster-a69601b5ee1fc2f2-v1-do-not-delete","cluster_user":"ec2-user","custom_awsbatchcli_package":"","custom_node_package":"","cw_logging_enabled":"true","default_user_home":"local","directory_service":{"domain_read_only_user":"","enabled":"false","generate_ssh_keys_for_users":"false"},"disable_sudo_access_for_default_user":"true","dns_domain":"{\\"Ref\\": \\"referencetoclusternameClusterDNSDomain8D0872E1Ref\\"}","ebs_shared_dirs":"","efs_encryption_in_transits":"","efs_fs_ids":"","efs_iam_authorizations":"","efs_shared_dirs":"","efs_access_point_ids":"","enable_intel_hpc_platform":"false","ephemeral_dir":"/scratch","fsx_dns_names":"","fsx_fs_ids":"","fsx_fs_types":"","fsx_mount_names":"","fsx_shared_dirs":"","fsx_volume_junction_paths":"","head_node_private_ip":"{\\"Ref\\": \\"referencetoclusternameHeadNodeENI6497A502PrimaryPrivateIpAddress\\"}","hosted_zone":"{\\"Ref\\": \\"referencetoclusternameRoute53HostedZone2388733DRef\\"}","dcv_enabled":"login_node","dcv_port":8444,"log_group_name":"/aws/parallelcluster/clustername-202401151530","log_rotation_enabled":"true","pool_name":"login","node_type":"LoginNode","proxy":"NONE","raid_shared_dir":"","raid_type":"","region":"us-east-1","scheduler":"slurm","shared_storage_type":"ebs","stack_arn":"{\\"Ref\\": \\"AWS::StackId\\"}","stack_name":"clustername","use_private_hostname":"false","launch_template_id":"LoginNodeLaunchTemplate2736fab291f04e69"}}\n'}, {'path': '/tmp/extra.json', 'permissions': '0644', 'owner': 'root:root', 'content': '{}\n'}, {'path': '/tmp/bootstrap.sh', 'permissions': '0744', 'owner': 'root:root', 'content': '#!/bin/bash -x\n\nfunction error_exit\n{\n echo "Bootstrap failed with error: $1"\n}\n'}]
44+
),
45+
("user_data_2.txt",
46+
[{'content': '{"cluster":{"base_os":"alinux2"}}\n', 'owner': 'root:root', 'path': '/tmp/dna.json', 'permissions': '0644'}, {'content': '{"cluster": {"nvidia": {"enabled": "yes" }, "is_official_ami_build": "true"}}\n', 'owner': 'root:root', 'path': '/tmp/extra.json', 'permissions': '0644'}, {'content': '#!/bin/bash -x\n\necho "Bootstrap failed with error: $1"\n', 'owner': 'root:root', 'path': '/tmp/bootstrap.sh', 'permissions': '0744'}]
47+
),
48+
("",None),
49+
],
50+
)
51+
def test_get_write_directives_section(mime_user_data_file, write_section, test_datadir):
52+
input_mime_user_data = None
53+
if mime_user_data_file:
54+
with open(os.path.join(test_datadir, mime_user_data_file), 'r') as file:
55+
input_mime_user_data = file.read().strip()
56+
57+
assert_that(get_write_directives_section(input_mime_user_data)).is_equal_to(write_section)
58+
59+
60+
@pytest.mark.parametrize(
61+
("error", "proxy", "port"),
62+
[
63+
(True, "myproxy.com", "8080"),
64+
(False, "", "")
65+
]
66+
)
67+
def test_parse_proxy_config(error, proxy, port):
68+
mock_config = MagicMock(return_value = error)
69+
mock_config.get.side_effect = [ proxy, port ]
70+
expected_op = {'https': proxy+':'+ port }
71+
with patch('configparser.RawConfigParser', return_value=mock_config):
72+
assert_that(parse_proxy_config().proxies).is_equal_to(expected_op)
73+
74+
75+
def ec2_describe_launch_template_versions_mock(response, lt_id, lt_version):
76+
e2_client = boto3.client('ec2', region_name='us-east-1')
77+
stubber = Stubber(e2_client)
78+
stubber.add_response('describe_launch_template_versions', response, {
79+
'LaunchTemplateId': lt_id,
80+
'Versions': [lt_version]
81+
})
82+
stubber.activate()
83+
return e2_client, stubber
84+
85+
86+
@pytest.mark.parametrize(
87+
('expected_user_data' ),
88+
[
89+
("#!/bin/bash\necho 'Test'"),
90+
("")
91+
92+
],
93+
)
94+
def test_get_user_data(expected_user_data):
95+
lt_id, lt_version = "lt-12345678901234567", "1"
96+
ec2_response = {
97+
"LaunchTemplateVersions": [{
98+
"LaunchTemplateData": {
99+
"UserData": b64encode(expected_user_data.encode()).decode('utf-8')
100+
}
101+
}]
102+
}
103+
104+
ec2_client, stubber = ec2_describe_launch_template_versions_mock(ec2_response, lt_id, lt_version)
105+
106+
with patch('boto3.client') as mock_client:
107+
mock_client.return_value = ec2_client
108+
with stubber:
109+
assert_that(get_user_data(lt_id, lt_version, 'us-east-1')).is_equal_to(expected_user_data)
110+
stubber.deactivate()

‎test/unit/cfn_hup_configuration/test_share_compute_fleet_dna/test_get_write_directives_section/user_data_1.txt‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
Content-Type: multipart/mixed; boundary="==BOUNDARY=="
2+
MIME-Version: 1.0
3+
4+
--==BOUNDARY==
5+
Content-Type: text/cloud-boothook; charset="us-ascii"
6+
MIME-Version: 1.0
7+
8+
#!/bin/bash -x
9+
10+
which dnf 2>/dev/null; dnf=$?
11+
which yum 2>/dev/null; yum=$?
12+
13+
--==BOUNDARY==
14+
Content-Type: text/cloud-config; charset=us-ascii
15+
MIME-Version: 1.0
16+
17+
bootcmd:
18+
19+
output:
20+
all: "| tee -a /var/log/cloud-init-output.log | logger -t user-data -s 2>/dev/ttyS0"
21+
write_files:
22+
- path: /tmp/dna.json
23+
permissions: '0644'
24+
owner: root:root
25+
content: |
26+
{"cluster":{"base_os":"alinux2","cluster_config_s3_key":"parallelcluster/clusters/dummy-cluster-randomstring123/configs/cluster-config-with-implied-values.yaml","cluster_config_version":"","cluster_name":"clustername","cluster_s3_bucket":"parallelcluster-a69601b5ee1fc2f2-v1-do-not-delete","cluster_user":"ec2-user","custom_awsbatchcli_package":"","custom_node_package":"","cw_logging_enabled":"true","default_user_home":"local","directory_service":{"domain_read_only_user":"","enabled":"false","generate_ssh_keys_for_users":"false"},"disable_sudo_access_for_default_user":"true","dns_domain":"{\"Ref\": \"referencetoclusternameClusterDNSDomain8D0872E1Ref\"}","ebs_shared_dirs":"","efs_encryption_in_transits":"","efs_fs_ids":"","efs_iam_authorizations":"","efs_shared_dirs":"","efs_access_point_ids":"","enable_intel_hpc_platform":"false","ephemeral_dir":"/scratch","fsx_dns_names":"","fsx_fs_ids":"","fsx_fs_types":"","fsx_mount_names":"","fsx_shared_dirs":"","fsx_volume_junction_paths":"","head_node_private_ip":"{\"Ref\": \"referencetoclusternameHeadNodeENI6497A502PrimaryPrivateIpAddress\"}","hosted_zone":"{\"Ref\": \"referencetoclusternameRoute53HostedZone2388733DRef\"}","dcv_enabled":"login_node","dcv_port":8444,"log_group_name":"/aws/parallelcluster/clustername-202401151530","log_rotation_enabled":"true","pool_name":"login","node_type":"LoginNode","proxy":"NONE","raid_shared_dir":"","raid_type":"","region":"us-east-1","scheduler":"slurm","shared_storage_type":"ebs","stack_arn":"{\"Ref\": \"AWS::StackId\"}","stack_name":"clustername","use_private_hostname":"false","launch_template_id":"LoginNodeLaunchTemplate2736fab291f04e69"}}
27+
- path: /tmp/extra.json
28+
permissions: '0644'
29+
owner: root:root
30+
content: |
31+
{}
32+
- path: /tmp/bootstrap.sh
33+
permissions: '0744'
34+
owner: root:root
35+
content: |
36+
#!/bin/bash -x
37+
38+
function error_exit
39+
{
40+
echo "Bootstrap failed with error: $1"
41+
}
42+
43+
--==BOUNDARY==
44+
Content-Type: text/x-shellscript; charset="us-ascii"
45+
MIME-Version: 1.0
46+
47+
#!/bin/bash -x
48+
49+
function error_exit
50+
{
51+
exit 1
52+
}
53+
--==BOUNDARY==

‎test/unit/cfn_hup_configuration/test_share_compute_fleet_dna/test_get_write_directives_section/user_data_2.txt‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
Content-Type: multipart/mixed; boundary="==BOUNDARY=="
2+
MIME-Version: 1.0
3+
4+
--==BOUNDARY==
5+
Content-Type: text/cloud-boothook; charset="us-ascii"
6+
MIME-Version: 1.0
7+
8+
#!/bin/bash -x
9+
10+
which dnf 2>/dev/null; dnf=$?
11+
which yum 2>/dev/null; yum=$?
12+
13+
--==BOUNDARY==
14+
Content-Type: text/cloud-config; charset=us-ascii
15+
MIME-Version: 1.0
16+
17+
write_files:
18+
- path: /tmp/dna.json
19+
permissions: '0644'
20+
owner: root:root
21+
content: |
22+
{"cluster":{"base_os":"alinux2"}}
23+
- path: /tmp/extra.json
24+
permissions: '0644'
25+
owner: root:root
26+
content: |
27+
{"cluster": {"nvidia": {"enabled": "yes" }, "is_official_ami_build": "true"}}
28+
- path: /tmp/bootstrap.sh
29+
permissions: '0744'
30+
owner: root:root
31+
content: |
32+
#!/bin/bash -x
33+
34+
echo "Bootstrap failed with error: $1"
35+
36+
--==BOUNDARY==
37+
Content-Type: text/x-shellscript; charset="us-ascii"
38+
MIME-Version: 1.0
39+
40+
#!/bin/bash -x
41+
function error_exit
42+
{
43+
exit 1
44+
}
45+
46+
--==BOUNDARY==

0 commit comments

Comments
 (0)