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,97 @@ 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" ,
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+ {"path" : "/tmp/extra.json" , "permissions" : "0644" , "owner" : "root:root" , "content" : "{}\n " },
90+ {
91+ "path" : "/tmp/bootstrap.sh" ,
92+ "permissions" : "0744" ,
93+ "owner" : "root:root" ,
94+ "content" : '#!/bin/bash -x\n \n function error_exit\n {\n echo "Bootstrap failed"\n }\n ' ,
95+ },
96+ ],
97+ ),
98+ (
99+ "user_data_2.txt" ,
100+ [
101+ {
102+ "content" : '{"cluster":{"base_os":"alinux2"}}\n ' ,
103+ "owner" : "root:root" ,
104+ "path" : "/tmp/dna.json" ,
105+ "permissions" : "0644" ,
106+ },
107+ {
108+ "content" : '{"cluster": {"nvidia": {"enabled": "yes" }, "is_official_ami_build": "true"}}\n ' ,
109+ "owner" : "root:root" ,
110+ "path" : "/tmp/extra.json" ,
111+ "permissions" : "0644" ,
112+ },
113+ {
114+ "content" : '#!/bin/bash -x\n \n echo "Bootstrap failed with error: $1"\n ' ,
115+ "owner" : "root:root" ,
116+ "path" : "/tmp/bootstrap.sh" ,
117+ "permissions" : "0744" ,
118+ },
119+ ],
120+ ),
121+ ("" , None ),
49122 ],
50123)
51124def test_get_write_directives_section (mime_user_data_file , write_section , test_datadir ):
52125 input_mime_user_data = None
53126 if mime_user_data_file :
54- with open (os .path .join (test_datadir , mime_user_data_file ), 'r' ) as file :
127+ with open (os .path .join (test_datadir , mime_user_data_file ), "r" ) as file :
55128 input_mime_user_data = file .read ().strip ()
56129
57130 assert_that (get_write_directives_section (input_mime_user_data )).is_equal_to (write_section )
58131
59132
60- @pytest .mark .parametrize (
61- ("error" , "proxy" , "port" ),
62- [
63- (True , "myproxy.com" , "8080" ),
64- (False , "" , "" )
65- ]
66- )
133+ @pytest .mark .parametrize (("error" , "proxy" , "port" ), [(True , "myproxy.com" , "8080" ), (False , "" , "" )])
67134def 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 ):
135+ mock_config = MagicMock (return_value = error )
136+ mock_config .get .side_effect = [proxy , port ]
137+ expected_op = {" https" proxy + ":" + port }
138+ with patch (" configparser.RawConfigParser" return_value = mock_config ):
72139 assert_that (parse_proxy_config ().proxies ).is_equal_to (expected_op )
73140
74141
75142def ec2_describe_launch_template_versions_mock (response , lt_id , lt_version ):
76- e2_client = boto3 .client (' ec2' region_name = ' us-east-1'
143+ e2_client = boto3 .client (" ec2" region_name = " us-east-1"
77144 stubber = Stubber (e2_client )
78- stubber .add_response ('describe_launch_template_versions' , response , {
79- 'LaunchTemplateId' : lt_id ,
80- 'Versions' : [lt_version ]
81- })
145+ stubber .add_response (
146+ "describe_launch_template_versions" , response , {"LaunchTemplateId" : lt_id , "Versions" : [lt_version ]}
147+ )
82148 stubber .activate ()
83149 return e2_client , stubber
84150
85151
86152@pytest .mark .parametrize (
87- ('expected_user_data' ),
88- [
89- ("#!/bin/bash\n echo 'Test'" ),
90- ("" )
91-
92- ],
153+ ("expected_user_data" ),
154+ [("#!/bin/bash\n echo 'Test'" ), ("" )],
93155)
94156def test_get_user_data (expected_user_data ):
95157 lt_id , lt_version = "lt-12345678901234567" , "1"
96158 ec2_response = {
97- "LaunchTemplateVersions" : [{
98- "LaunchTemplateData" : {
99- "UserData" : b64encode (expected_user_data .encode ()).decode ('utf-8' )
100- }
101- }]
102- }
159+ "LaunchTemplateVersions" : [
160+ {"LaunchTemplateData" : {"UserData" : b64encode (expected_user_data .encode ()).decode ("utf-8" )}}
161+ ]
162+ }
103163
104164 ec2_client , stubber = ec2_describe_launch_template_versions_mock (ec2_response , lt_id , lt_version )
105165
106- with patch (' boto3.client' as mock_client :
166+ with patch (" boto3.client" as mock_client :
107167 mock_client .return_value = ec2_client
108168 with stubber :
109- assert_that (get_user_data (lt_id , lt_version , ' us-east-1' is_equal_to (expected_user_data )
110- stubber .deactivate ()
169+ assert_that (get_user_data (lt_id , lt_version , " us-east-1" is_equal_to (expected_user_data )
170+ stubber .deactivate ()
0 commit comments