5
5
from models_library .docker import DockerGenericTag
6
6
from types_aiobotocore_ec2 import EC2Client
7
7
from types_aiobotocore_ec2 .literals import InstanceStateNameType , InstanceTypeType
8
- from types_aiobotocore_ec2 .type_defs import FilterTypeDef , InstanceTypeDef , TagTypeDef
8
+ from types_aiobotocore_ec2 .type_defs import (
9
+ FilterTypeDef ,
10
+ InstanceTypeDef ,
11
+ ReservationTypeDef ,
12
+ TagTypeDef ,
13
+ )
9
14
10
15
11
16
async def assert_autoscaled_computational_ec2_instances (
@@ -43,6 +48,7 @@ async def assert_autoscaled_dynamic_ec2_instances(
43
48
expected_additional_tag_keys : list [str ],
44
49
instance_filters : Sequence [FilterTypeDef ] | None ,
45
50
expected_user_data : list [str ] | None = None ,
51
+ check_reservation_index : int | None = None ,
46
52
) -> list [InstanceTypeDef ]:
47
53
if expected_user_data is None :
48
54
expected_user_data = ["docker swarm join" ]
@@ -59,6 +65,7 @@ async def assert_autoscaled_dynamic_ec2_instances(
59
65
],
60
66
expected_user_data = expected_user_data ,
61
67
instance_filters = instance_filters ,
68
+ check_reservation_index = check_reservation_index ,
62
69
)
63
70
64
71
@@ -72,6 +79,7 @@ async def assert_autoscaled_dynamic_warm_pools_ec2_instances(
72
79
expected_additional_tag_keys : list [str ],
73
80
expected_pre_pulled_images : list [DockerGenericTag ] | None ,
74
81
instance_filters : Sequence [FilterTypeDef ] | None ,
82
+ check_reservation_index : int | None = None ,
75
83
) -> list [InstanceTypeDef ]:
76
84
return await assert_ec2_instances (
77
85
ec2_client ,
@@ -88,9 +96,80 @@ async def assert_autoscaled_dynamic_warm_pools_ec2_instances(
88
96
expected_pre_pulled_images = expected_pre_pulled_images ,
89
97
expected_user_data = [],
90
98
instance_filters = instance_filters ,
99
+ check_reservation_index = check_reservation_index ,
91
100
)
92
101
93
102
103
+ async def _assert_reservation (
104
+ ec2_client : EC2Client ,
105
+ reservation : ReservationTypeDef ,
106
+ * ,
107
+ expected_num_instances : int ,
108
+ expected_instance_type : InstanceTypeType ,
109
+ expected_instance_state : InstanceStateNameType ,
110
+ expected_instance_tag_keys : list [str ],
111
+ expected_user_data : list [str ],
112
+ expected_pre_pulled_images : list [DockerGenericTag ] | None ,
113
+ ) -> list [InstanceTypeDef ]:
114
+ list_instances : list [InstanceTypeDef ] = []
115
+ assert "Instances" in reservation
116
+ assert (
117
+ len (reservation ["Instances" ]) == expected_num_instances
118
+ ), f"expected { expected_num_instances } , found { len (reservation ['Instances' ])} "
119
+ for instance in reservation ["Instances" ]:
120
+ assert "InstanceType" in instance
121
+ assert instance ["InstanceType" ] == expected_instance_type
122
+ assert "Tags" in instance
123
+ assert instance ["Tags" ]
124
+ expected_tag_keys = {
125
+ * expected_instance_tag_keys ,
126
+ "io.simcore.autoscaling.version" ,
127
+ "Name" ,
128
+ }
129
+ instance_tag_keys = {tag ["Key" ] for tag in instance ["Tags" ] if "Key" in tag }
130
+ assert instance_tag_keys == expected_tag_keys
131
+
132
+ if expected_pre_pulled_images is None :
133
+ assert "io.simcore.autoscaling.pre_pulled_images" not in instance_tag_keys
134
+ else :
135
+ assert "io.simcore.autoscaling.pre_pulled_images" in instance_tag_keys
136
+
137
+ def _by_pre_pull_image (ec2_tag : TagTypeDef ) -> bool :
138
+ assert "Key" in ec2_tag
139
+ return ec2_tag ["Key" ] == "io.simcore.autoscaling.pre_pulled_images"
140
+
141
+ instance_pre_pulled_images_aws_tag = next (
142
+ iter (filter (_by_pre_pull_image , instance ["Tags" ]))
143
+ )
144
+ assert "Value" in instance_pre_pulled_images_aws_tag
145
+ assert (
146
+ instance_pre_pulled_images_aws_tag ["Value" ]
147
+ == f"{ json_dumps (expected_pre_pulled_images )} "
148
+ )
149
+
150
+ assert "PrivateDnsName" in instance
151
+ instance_private_dns_name = instance ["PrivateDnsName" ]
152
+ if expected_instance_state not in ["terminated" ]:
153
+ # NOTE: moto behaves here differently than AWS by still returning an IP which does not really make sense
154
+ assert instance_private_dns_name .endswith (".ec2.internal" )
155
+ assert "State" in instance
156
+ state = instance ["State" ]
157
+ assert "Name" in state
158
+ assert state ["Name" ] == expected_instance_state
159
+
160
+ assert "InstanceId" in instance
161
+ user_data = await ec2_client .describe_instance_attribute (
162
+ Attribute = "userData" , InstanceId = instance ["InstanceId" ]
163
+ )
164
+ assert "UserData" in user_data
165
+ assert "Value" in user_data ["UserData" ]
166
+ user_data = base64 .b64decode (user_data ["UserData" ]["Value" ]).decode ()
167
+ for user_data_string in expected_user_data :
168
+ assert user_data .count (user_data_string ) == 1
169
+ list_instances .append (instance )
170
+ return list_instances
171
+
172
+
94
173
async def assert_ec2_instances (
95
174
ec2_client : EC2Client ,
96
175
* ,
@@ -102,66 +181,35 @@ async def assert_ec2_instances(
102
181
expected_user_data : list [str ],
103
182
expected_pre_pulled_images : list [DockerGenericTag ] | None = None ,
104
183
instance_filters : Sequence [FilterTypeDef ] | None = None ,
184
+ check_reservation_index : int | None = None ,
105
185
) -> list [InstanceTypeDef ]:
106
- list_instances : list [InstanceTypeDef ] = []
107
186
all_instances = await ec2_client .describe_instances (Filters = instance_filters or [])
108
187
assert len (all_instances ["Reservations" ]) == expected_num_reservations
188
+ if check_reservation_index is not None :
189
+ assert check_reservation_index < len (all_instances ["Reservations" ])
190
+ reservation = all_instances ["Reservations" ][check_reservation_index ]
191
+ return await _assert_reservation (
192
+ ec2_client ,
193
+ reservation ,
194
+ expected_num_instances = expected_num_instances ,
195
+ expected_instance_type = expected_instance_type ,
196
+ expected_instance_state = expected_instance_state ,
197
+ expected_instance_tag_keys = expected_instance_tag_keys ,
198
+ expected_user_data = expected_user_data ,
199
+ expected_pre_pulled_images = expected_pre_pulled_images ,
200
+ )
201
+ list_instances : list [InstanceTypeDef ] = []
109
202
for reservation in all_instances ["Reservations" ]:
110
- assert "Instances" in reservation
111
- assert (
112
- len (reservation ["Instances" ]) == expected_num_instances
113
- ), f"expected { expected_num_instances } , found { len (reservation ['Instances' ])} "
114
- for instance in reservation ["Instances" ]:
115
- assert "InstanceType" in instance
116
- assert instance ["InstanceType" ] == expected_instance_type
117
- assert "Tags" in instance
118
- assert instance ["Tags" ]
119
- expected_tag_keys = {
120
- * expected_instance_tag_keys ,
121
- "io.simcore.autoscaling.version" ,
122
- "Name" ,
123
- }
124
- instance_tag_keys = {tag ["Key" ] for tag in instance ["Tags" ] if "Key" in tag }
125
- assert instance_tag_keys == expected_tag_keys
126
-
127
- if expected_pre_pulled_images is None :
128
- assert (
129
- "io.simcore.autoscaling.pre_pulled_images" not in instance_tag_keys
130
- )
131
- else :
132
- assert "io.simcore.autoscaling.pre_pulled_images" in instance_tag_keys
133
-
134
- def _by_pre_pull_image (ec2_tag : TagTypeDef ) -> bool :
135
- assert "Key" in ec2_tag
136
- return ec2_tag ["Key" ] == "io.simcore.autoscaling.pre_pulled_images"
137
-
138
- instance_pre_pulled_images_aws_tag = next (
139
- iter (filter (_by_pre_pull_image , instance ["Tags" ]))
140
- )
141
- assert "Value" in instance_pre_pulled_images_aws_tag
142
- assert (
143
- instance_pre_pulled_images_aws_tag ["Value" ]
144
- == f"{ json_dumps (expected_pre_pulled_images )} "
145
- )
146
-
147
- assert "PrivateDnsName" in instance
148
- instance_private_dns_name = instance ["PrivateDnsName" ]
149
- if expected_instance_state not in ["terminated" ]:
150
- # NOTE: moto behaves here differently than AWS by still returning an IP which does not really make sense
151
- assert instance_private_dns_name .endswith (".ec2.internal" )
152
- assert "State" in instance
153
- state = instance ["State" ]
154
- assert "Name" in state
155
- assert state ["Name" ] == expected_instance_state
156
-
157
- assert "InstanceId" in instance
158
- user_data = await ec2_client .describe_instance_attribute (
159
- Attribute = "userData" , InstanceId = instance ["InstanceId" ]
203
+ list_instances .extend (
204
+ await _assert_reservation (
205
+ ec2_client ,
206
+ reservation ,
207
+ expected_num_instances = expected_num_instances ,
208
+ expected_instance_type = expected_instance_type ,
209
+ expected_instance_state = expected_instance_state ,
210
+ expected_instance_tag_keys = expected_instance_tag_keys ,
211
+ expected_user_data = expected_user_data ,
212
+ expected_pre_pulled_images = expected_pre_pulled_images ,
160
213
)
161
- assert "UserData" in user_data
162
- assert "Value" in user_data ["UserData" ]
163
- user_data = base64 .b64decode (user_data ["UserData" ]["Value" ]).decode ()
164
- for user_data_string in expected_user_data :
165
- assert user_data .count (user_data_string ) == 1
166
- list_instances .append (instance )
214
+ )
167
215
return list_instances
0 commit comments