88# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
99# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
1010# limitations under the License.
11+ import json
1112import os
13+ from base64 import b64encode
1214from unittest .mock import MagicMock , patch
15+
1316import boto3
1417import pytest
15- import json
1618from assertpy import assert_that
1719from 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+ from share_compute_fleet_dna import (
21+ get_compute_launch_template_ids ,
22+ get_user_data ,
23+ get_write_directives_section ,
24+ parse_proxy_config ,
25+ )
2026
2127
2228@pytest .mark .parametrize (
2329 ("launch_template_config_content" , "errors" ),
2430 [
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-
31+ (
32+ """
33+ {
34+ "Queues": {
35+ "queue-0": {
36+ "ComputeResources": {
37+ "compute-resource-1": {
38+ "LaunchTemplate": {
39+ "Version": "1",
40+ "Id": "lt-037123456747c3bc5"
41+ }
42+ },
43+ "compute-resource-2": {
44+ "LaunchTemplate": {
45+ "Version": "1",
46+ "Id": "lt-0fcecb59a3721c0b3"
47+ }
48+ },
49+ "compute-resource-0": {
50+ "LaunchTemplate": {
51+ "Version": "1",
52+ "Id": "lt-12345678901234567"
53+ }
54+ }
55+ }
56+ }
57+ }
58+ }
59+ """ ,
60+ False ,
61+ ),
62+ ('{"Queues":{"queue-0":}}}' , True ),
2863 ],
2964)
3065def test_get_compute_launch_template_ids (mocker , launch_template_config_content , errors ):
@@ -39,72 +74,102 @@ def test_get_compute_launch_template_ids(mocker, launch_template_config_content,
3974@pytest .mark .parametrize (
4075 ("mime_user_data_file" , "write_section" ),
4176 [
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 \n function 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 \n echo "Bootstrap failed with error: $1"\n ' , 'owner' : 'root:root' , 'path' : '/tmp/bootstrap.sh' , 'permissions' : '0744' }]
47- ),
48- ("" ,None ),
77+ (
78+ "user_data_1.txt" ,
79+ [
80+ {
81+ "path" : "/tmp/dna.json" , # nosec B108
82+ "permissions" : "0644" ,
83+ "owner" : "root:root" ,
84+ "content" : '{"cluster":{"base_os":"alinux2","cluster_name":"clustername",'
85+ '"directory_service":{"domain_read_only_user":"","enabled":"false",'
86+ '"generate_ssh_keys_for_users":"false"},'
87+ '"launch_template_id":"LoginNodeLaunchTemplate2736fab291f04e69"}}\n ' ,
88+ },
89+ {
90+ "path" : "/tmp/extra.json" , # nosec B108
91+ "permissions" : "0644" ,
92+ "owner" : "root:root" ,
93+ "content" : "{}\n " ,
94+ },
95+ {
96+ "path" : "/tmp/bootstrap.sh" , # nosec B108
97+ "permissions" : "0744" ,
98+ "owner" : "root:root" ,
99+ "content" : '#!/bin/bash -x\n \n function error_exit\n {\n echo "Bootstrap failed"\n }\n ' ,
100+ },
101+ ],
102+ ),
103+ (
104+ "user_data_2.txt" ,
105+ [
106+ {
107+ "content" : '{"cluster":{"base_os":"alinux2"}}\n ' ,
108+ "owner" : "root:root" ,
109+ "path" : "/tmp/dna.json" , # nosec B108
110+ "permissions" : "0644" ,
111+ },
112+ {
113+ "content" : '{"cluster": {"nvidia": {"enabled": "yes" }, "is_official_ami_build": "true"}}\n ' ,
114+ "owner" : "root:root" ,
115+ "path" : "/tmp/extra.json" , # nosec B108
116+ "permissions" : "0644" ,
117+ },
118+ {
119+ "content" : '#!/bin/bash -x\n \n echo "Bootstrap failed with error: $1"\n ' ,
120+ "owner" : "root:root" ,
121+ "path" : "/tmp/bootstrap.sh" , # nosec B108
122+ "permissions" : "0744" ,
123+ },
124+ ],
125+ ),
126+ ("" , None ),
49127 ],
50128)
51129def test_get_write_directives_section (mime_user_data_file , write_section , test_datadir ):
52130 input_mime_user_data = None
53131 if mime_user_data_file :
54- with open (os .path .join (test_datadir , mime_user_data_file ), 'r' ) as file :
132+ with open (os .path .join (test_datadir , mime_user_data_file ), "r" ) as file :
55133 input_mime_user_data = file .read ().strip ()
56134
57135 assert_that (get_write_directives_section (input_mime_user_data )).is_equal_to (write_section )
58136
59137
60- @pytest .mark .parametrize (
61- ("error" , "proxy" , "port" ),
62- [
63- (True , "myproxy.com" , "8080" ),
64- (False , "" , "" )
65- ]
66- )
138+ @pytest .mark .parametrize (("error" , "proxy" , "port" ), [(True , "myproxy.com" , "8080" ), (False , "" , "" )])
67139def 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 ):
140+ mock_config = MagicMock (return_value = error )
141+ mock_config .get .side_effect = [proxy , port ]
142+ expected_op = {" https" proxy + ":" + port }
143+ with patch (" configparser.RawConfigParser" return_value = mock_config ):
72144 assert_that (parse_proxy_config ().proxies ).is_equal_to (expected_op )
73145
74146
75147def ec2_describe_launch_template_versions_mock (response , lt_id , lt_version ):
76- e2_client = boto3 .client (' ec2' region_name = ' us-east-1'
148+ e2_client = boto3 .client (" ec2" region_name = " us-east-1"
77149 stubber = Stubber (e2_client )
78- stubber .add_response ('describe_launch_template_versions' , response , {
79- 'LaunchTemplateId' : lt_id ,
80- 'Versions' : [lt_version ]
81- })
150+ stubber .add_response (
151+ "describe_launch_template_versions" , response , {"LaunchTemplateId" : lt_id , "Versions" : [lt_version ]}
152+ )
82153 stubber .activate ()
83154 return e2_client , stubber
84155
85156
86157@pytest .mark .parametrize (
87- ('expected_user_data' ),
88- [
89- ("#!/bin/bash\n echo 'Test'" ),
90- ("" )
91-
92- ],
158+ ("expected_user_data" ),
159+ [("#!/bin/bash\n echo 'Test'" ), ("" )],
93160)
94161def test_get_user_data (expected_user_data ):
95162 lt_id , lt_version = "lt-12345678901234567" , "1"
96163 ec2_response = {
97- "LaunchTemplateVersions" : [{
98- "LaunchTemplateData" : {
99- "UserData" : b64encode (expected_user_data .encode ()).decode ('utf-8' )
100- }
101- }]
102- }
164+ "LaunchTemplateVersions" : [
165+ {"LaunchTemplateData" : {"UserData" : b64encode (expected_user_data .encode ()).decode ("utf-8" )}}
166+ ]
167+ }
103168
104169 ec2_client , stubber = ec2_describe_launch_template_versions_mock (ec2_response , lt_id , lt_version )
105170
106- with patch (' boto3.client' as mock_client :
171+ with patch (" boto3.client" as mock_client :
107172 mock_client .return_value = ec2_client
108173 with stubber :
109- assert_that (get_user_data (lt_id , lt_version , ' us-east-1' is_equal_to (expected_user_data )
110- stubber .deactivate ()
174+ assert_that (get_user_data (lt_id , lt_version , " us-east-1" is_equal_to (expected_user_data )
175+ stubber .deactivate ()
0 commit comments