Skip to content

Commit 9287307

Browse files
chenwanychenwany
committed
Add unit test for runinstance override
Signed-off-by: chenwany <[email protected]>
1 parent 5bd43ce commit 9287307

File tree

1 file changed

+289
-1
lines changed

1 file changed

+289
-1
lines changed

‎tests/slurm_plugin/test_common.py‎

Lines changed: 289 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,20 @@ def instance_manager(self, mocker):
6363
dns_domain="dns.domain",
6464
use_private_hostname=False,
6565
instance_name_type_mapping={"c5xlarge": "c5.xlarge"},
66+
run_instances_overrides={
67+
"queue3": {
68+
"p4d.24xlarge": {
69+
"CapacityReservationSpecification": {
70+
"CapacityReservationTarget": {"CapacityReservationId": "cr-123"}
71+
}
72+
},
73+
"c5.xlarge": {
74+
"CapacityReservationSpecification": {
75+
"CapacityReservationTarget": {"CapacityReservationId": "cr-456"}
76+
}
77+
},
78+
},
79+
},
6680
)
6781
mocker.patch.object(instance_manager, "_table")
6882
return instance_manager
@@ -603,8 +617,148 @@ def instance_manager(self, mocker):
603617
)
604618
],
605619
),
620+
# override_runinstances
621+
(
622+
{
623+
"queue3": {
624+
"c5.xlarge": ["queue3-st-c5xlarge-2"],
625+
"c5.2xlarge": ["queue3-st-c52xlarge-1"],
626+
"p4d.24xlarge": ["queue3-st-p4d24xlarge-1"],
627+
},
628+
"queue2": {"c5.xlarge": ["queue2-st-c5xlarge-1", "queue2-dy-c5xlarge-1"]},
629+
},
630+
10,
631+
True,
632+
False,
633+
[
634+
MockedBoto3Request(
635+
method="run_instances",
636+
response={
637+
"Instances": [
638+
{
639+
"InstanceId": "i-12345",
640+
"InstanceType": "c5.xlarge",
641+
"PrivateIpAddress": "ip.1.0.0.1",
642+
"PrivateDnsName": "ip-1-0-0-1",
643+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
644+
}
645+
]
646+
},
647+
expected_params={
648+
"MinCount": 1,
649+
"MaxCount": 1,
650+
"LaunchTemplate": {"LaunchTemplateName": "hit-queue3-c5.xlarge", "Version": "$Latest"},
651+
"CapacityReservationSpecification": {
652+
"CapacityReservationTarget": {"CapacityReservationId": "cr-456"}
653+
},
654+
},
655+
),
656+
MockedBoto3Request(
657+
method="run_instances",
658+
response={
659+
"Instances": [
660+
{
661+
"InstanceId": "i-23456",
662+
"InstanceType": "c5.2xlarge",
663+
"PrivateIpAddress": "ip.1.0.0.2",
664+
"PrivateDnsName": "ip-1-0-0-2",
665+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
666+
}
667+
]
668+
},
669+
expected_params={
670+
"MinCount": 1,
671+
"MaxCount": 1,
672+
"LaunchTemplate": {"LaunchTemplateName": "hit-queue3-c5.2xlarge", "Version": "$Latest"},
673+
},
674+
),
675+
MockedBoto3Request(
676+
method="run_instances",
677+
response={
678+
"Instances": [
679+
{
680+
"InstanceId": "i-12346",
681+
"InstanceType": "p4d.24xlarge",
682+
"PrivateIpAddress": "ip.1.0.0.4",
683+
"PrivateDnsName": "ip-1-0-0-4",
684+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
685+
}
686+
]
687+
},
688+
expected_params={
689+
"MinCount": 1,
690+
"MaxCount": 1,
691+
"LaunchTemplate": {"LaunchTemplateName": "hit-queue3-p4d.24xlarge", "Version": "$Latest"},
692+
"CapacityReservationSpecification": {
693+
"CapacityReservationTarget": {"CapacityReservationId": "cr-123"}
694+
},
695+
},
696+
),
697+
MockedBoto3Request(
698+
method="run_instances",
699+
response={
700+
"Instances": [
701+
{
702+
"InstanceId": "i-34567",
703+
"InstanceType": "c5.xlarge",
704+
"PrivateIpAddress": "ip.1.0.0.3",
705+
"PrivateDnsName": "ip-1-0-0-3",
706+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
707+
},
708+
{
709+
"InstanceId": "i-45678",
710+
"InstanceType": "c5.xlarge",
711+
"PrivateIpAddress": "ip.1.0.0.4",
712+
"PrivateDnsName": "ip-1-0-0-4",
713+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
714+
},
715+
]
716+
},
717+
expected_params={
718+
"MinCount": 1,
719+
"MaxCount": 2,
720+
"LaunchTemplate": {"LaunchTemplateName": "hit-queue2-c5.xlarge", "Version": "$Latest"},
721+
},
722+
),
723+
],
724+
None,
725+
[
726+
call(
727+
["queue3-st-c5xlarge-2"],
728+
[EC2Instance("i-12345", "ip.1.0.0.1", "ip-1-0-0-1", datetime(2020, 1, 1, tzinfo=timezone.utc))],
729+
),
730+
call(
731+
["queue3-st-c52xlarge-1"],
732+
[EC2Instance("i-23456", "ip.1.0.0.2", "ip-1-0-0-2", datetime(2020, 1, 1, tzinfo=timezone.utc))],
733+
),
734+
call(
735+
["queue3-st-p4d24xlarge-1"],
736+
[EC2Instance("i-12346", "ip.1.0.0.4", "ip-1-0-0-4", datetime(2020, 1, 1, tzinfo=timezone.utc))],
737+
),
738+
call(
739+
["queue2-st-c5xlarge-1", "queue2-dy-c5xlarge-1"],
740+
[
741+
EC2Instance(
742+
"i-34567", "ip.1.0.0.3", "ip-1-0-0-3", datetime(2020, 1, 1, tzinfo=timezone.utc)
743+
),
744+
EC2Instance(
745+
"i-45678", "ip.1.0.0.4", "ip-1-0-0-4", datetime(2020, 1, 1, tzinfo=timezone.utc)
746+
),
747+
],
748+
),
749+
],
750+
),
751+
],
752+
ids=[
753+
"normal",
754+
"client_error",
755+
"no_update",
756+
"batch_size1",
757+
"batch_size2",
758+
"partial_launch",
759+
"all_or_nothing",
760+
"override_runinstances",
606761
],
607-
ids=["normal", "client_error", "no_update", "batch_size1", "batch_size2", "partial_launch", "all_or_nothing"],
608762
)
609763
def test_add_instances(
610764
self,
@@ -1309,6 +1463,140 @@ def test_get_cluster_instances(
13091463
result = instance_manager.get_cluster_instances(**mock_kwargs)
13101464
assert_that(result).is_equal_to(expected_parsed_result)
13111465

1466+
@pytest.mark.parametrize(
1467+
(
1468+
"current_batch_size",
1469+
"instance_type",
1470+
"all_or_nothing_batch",
1471+
"run_instances_overrides",
1472+
"mocked_boto3_request",
1473+
),
1474+
[
1475+
(
1476+
5,
1477+
"p4d.24xlarge",
1478+
False,
1479+
{},
1480+
[
1481+
MockedBoto3Request(
1482+
method="run_instances",
1483+
response={
1484+
"Instances": [
1485+
{
1486+
"InstanceId": "i-12345",
1487+
"InstanceType": "p4d.24xlarge",
1488+
"PrivateIpAddress": "ip.1.0.0.1",
1489+
"PrivateDnsName": "ip-1-0-0-1",
1490+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
1491+
}
1492+
]
1493+
},
1494+
expected_params={
1495+
"MinCount": 1,
1496+
"MaxCount": 5,
1497+
"LaunchTemplate": {
1498+
"LaunchTemplateName": "hit-queue-p4d.24xlarge",
1499+
"Version": "$Latest",
1500+
},
1501+
},
1502+
),
1503+
],
1504+
),
1505+
(
1506+
5,
1507+
"c5.xlarge",
1508+
True,
1509+
{},
1510+
[
1511+
MockedBoto3Request(
1512+
method="run_instances",
1513+
response={
1514+
"Instances": [
1515+
{
1516+
"InstanceId": "i-12345",
1517+
"InstanceType": "c5.xlarge",
1518+
"PrivateIpAddress": "ip.1.0.0.1",
1519+
"PrivateDnsName": "ip-1-0-0-1",
1520+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
1521+
}
1522+
]
1523+
},
1524+
expected_params={
1525+
"MinCount": 5,
1526+
"MaxCount": 5,
1527+
"LaunchTemplate": {
1528+
"LaunchTemplateName": "hit-queue-c5.xlarge",
1529+
"Version": "$Latest",
1530+
},
1531+
},
1532+
),
1533+
],
1534+
),
1535+
(
1536+
5,
1537+
"p4d.24xlarge",
1538+
False,
1539+
{
1540+
"CapacityReservationSpecification": {
1541+
"CapacityReservationTarget": {"CapacityReservationId": "cr-12345"}
1542+
}
1543+
},
1544+
[
1545+
MockedBoto3Request(
1546+
method="run_instances",
1547+
response={
1548+
"Instances": [
1549+
{
1550+
"InstanceId": "i-12345",
1551+
"InstanceType": "p4d.24xlarge",
1552+
"PrivateIpAddress": "ip.1.0.0.1",
1553+
"PrivateDnsName": "ip-1-0-0-1",
1554+
"LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
1555+
}
1556+
]
1557+
},
1558+
expected_params={
1559+
"MinCount": 1,
1560+
"MaxCount": 5,
1561+
"LaunchTemplate": {
1562+
"LaunchTemplateName": "hit-queue-p4d.24xlarge",
1563+
"Version": "$Latest",
1564+
},
1565+
"CapacityReservationSpecification": {
1566+
"CapacityReservationTarget": {"CapacityReservationId": "cr-12345"}
1567+
},
1568+
},
1569+
),
1570+
],
1571+
),
1572+
],
1573+
ids=["normal", "all_or_nothing_batch", "run_instances_overrides"],
1574+
)
1575+
def test_launch_ec2_instances(
1576+
self,
1577+
boto3_stubber,
1578+
current_batch_size,
1579+
instance_type,
1580+
all_or_nothing_batch,
1581+
run_instances_overrides,
1582+
mocked_boto3_request,
1583+
instance_manager,
1584+
caplog,
1585+
):
1586+
caplog.set_level(logging.INFO)
1587+
# patch boto3 call
1588+
boto3_stubber("ec2", mocked_boto3_request)
1589+
# run test
1590+
instance_manager._launch_ec2_instances(
1591+
"queue",
1592+
instance_type,
1593+
current_batch_size,
1594+
all_or_nothing_batch=all_or_nothing_batch,
1595+
run_instances_overrides=run_instances_overrides,
1596+
)
1597+
if run_instances_overrides:
1598+
assert_that(caplog.text).contains("Found RunInstances parameters override")
1599+
13121600

13131601
@pytest.mark.parametrize(
13141602
"initial_time, current_time, grace_time, expected_result",

0 commit comments

Comments
 (0)