Skip to content

Commit 621ff80

Browse files
sandereggsanderegg
committed
AI step 2
1 parent 9cec957 commit 621ff80

File tree

1 file changed

+302
-0
lines changed

1 file changed

+302
-0
lines changed

‎packages/aws-library/tests/test_ec2_client.py‎

Lines changed: 302 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,15 @@
77
from collections.abc import AsyncIterator, Callable
88
from dataclasses import fields
99
from typing import cast, get_args
10+
from unittest.mock import patch
1011

1112
import botocore.exceptions
1213
import pytest
1314
from aws_library.ec2._client import SimcoreEC2API
1415
from aws_library.ec2._errors import (
1516
EC2InstanceNotFoundError,
1617
EC2InstanceTypeInvalidError,
18+
EC2InsufficientCapacityError,
1719
EC2TooManyInstancesError,
1820
)
1921
from aws_library.ec2._models import (
@@ -617,3 +619,303 @@ async def test_launch_instances_multi_subnet_fallback(
617619
expected_tags=ec2_instance_config.tags,
618620
expected_state="running",
619621
)
622+
623+
624+
async def test_launch_instances_multi_subnet_fallback_with_invalid_subnets(
625+
simcore_ec2_api: SimcoreEC2API,
626+
ec2_client: EC2Client,
627+
fake_ec2_instance_type: EC2InstanceType,
628+
faker: Faker,
629+
aws_subnet_id: str,
630+
aws_security_group_id: str,
631+
aws_ami_id: str,
632+
):
633+
"""Test that launch_instances uses multiple valid subnets correctly."""
634+
await _assert_no_instances_in_ec2(ec2_client)
635+
636+
# Create additional valid subnets for testing
637+
vpc_id = (await ec2_client.describe_subnets(SubnetIds=[aws_subnet_id]))["Subnets"][
638+
0
639+
]["VpcId"]
640+
641+
# Create second subnet
642+
subnet2 = await ec2_client.create_subnet(CidrBlock="10.0.2.0/24", VpcId=vpc_id)
643+
subnet2_id = subnet2["Subnet"]["SubnetId"]
644+
645+
# Create third subnet
646+
subnet3 = await ec2_client.create_subnet(CidrBlock="10.0.3.0/24", VpcId=vpc_id)
647+
subnet3_id = subnet3["Subnet"]["SubnetId"]
648+
649+
try:
650+
# Create a config with multiple valid subnet IDs
651+
ec2_instance_config = EC2InstanceConfig(
652+
type=fake_ec2_instance_type,
653+
tags=faker.pydict(allowed_types=(str,)),
654+
startup_script=faker.pystr(),
655+
ami_id=aws_ami_id,
656+
key_name=faker.pystr(),
657+
security_group_ids=[aws_security_group_id],
658+
subnet_ids=[aws_subnet_id, subnet2_id, subnet3_id],
659+
iam_instance_profile="",
660+
)
661+
662+
# This should succeed using one of the valid subnets
663+
instances = await simcore_ec2_api.launch_instances(
664+
ec2_instance_config,
665+
min_number_of_instances=1,
666+
number_of_instances=1,
667+
)
668+
669+
# Verify that the instance was created in one of the configured subnets
670+
await _assert_instances_in_ec2(
671+
ec2_client,
672+
expected_num_reservations=1,
673+
expected_num_instances=1,
674+
expected_instance_type=ec2_instance_config.type,
675+
expected_tags=ec2_instance_config.tags,
676+
expected_state="running",
677+
)
678+
679+
# Verify the instance was created in one of the valid subnets
680+
instance_details = await ec2_client.describe_instances(
681+
InstanceIds=[instances[0].id]
682+
)
683+
instance = instance_details["Reservations"][0]["Instances"][0]
684+
assert instance["SubnetId"] in [aws_subnet_id, subnet2_id, subnet3_id]
685+
686+
finally:
687+
# Clean up additional subnets
688+
await ec2_client.delete_subnet(SubnetId=subnet2_id)
689+
await ec2_client.delete_subnet(SubnetId=subnet3_id)
690+
691+
692+
async def test_launch_instances_multi_subnet_configuration(
693+
simcore_ec2_api: SimcoreEC2API,
694+
ec2_client: EC2Client,
695+
fake_ec2_instance_type: EC2InstanceType,
696+
faker: Faker,
697+
aws_subnet_id: str,
698+
aws_security_group_id: str,
699+
aws_ami_id: str,
700+
):
701+
"""Test that launch_instances works with multiple valid subnets configuration."""
702+
await _assert_no_instances_in_ec2(ec2_client)
703+
704+
# Create additional valid subnets for testing
705+
vpc_id = (await ec2_client.describe_subnets(SubnetIds=[aws_subnet_id]))["Subnets"][
706+
0
707+
]["VpcId"]
708+
709+
# Create second subnet
710+
subnet2 = await ec2_client.create_subnet(CidrBlock="10.0.2.0/24", VpcId=vpc_id)
711+
subnet2_id = subnet2["Subnet"]["SubnetId"]
712+
713+
# Create third subnet
714+
subnet3 = await ec2_client.create_subnet(CidrBlock="10.0.3.0/24", VpcId=vpc_id)
715+
subnet3_id = subnet3["Subnet"]["SubnetId"]
716+
717+
try:
718+
# Create a config with multiple valid subnet IDs
719+
ec2_instance_config = EC2InstanceConfig(
720+
type=fake_ec2_instance_type,
721+
tags=faker.pydict(allowed_types=(str,)),
722+
startup_script=faker.pystr(),
723+
ami_id=aws_ami_id,
724+
key_name=faker.pystr(),
725+
security_group_ids=[aws_security_group_id],
726+
subnet_ids=[aws_subnet_id, subnet2_id, subnet3_id],
727+
iam_instance_profile="",
728+
)
729+
730+
# This should succeed using one of the valid subnets
731+
instances = await simcore_ec2_api.launch_instances(
732+
ec2_instance_config,
733+
min_number_of_instances=1,
734+
number_of_instances=1,
735+
)
736+
737+
# Verify that the instance was created in one of the configured subnets
738+
await _assert_instances_in_ec2(
739+
ec2_client,
740+
expected_num_reservations=1,
741+
expected_num_instances=1,
742+
expected_instance_type=ec2_instance_config.type,
743+
expected_tags=ec2_instance_config.tags,
744+
expected_state="running",
745+
)
746+
747+
# Verify the instance was created in one of the valid subnets
748+
instance_details = await ec2_client.describe_instances(
749+
InstanceIds=[instances[0].id]
750+
)
751+
instance = instance_details["Reservations"][0]["Instances"][0]
752+
assert instance["SubnetId"] in [aws_subnet_id, subnet2_id, subnet3_id]
753+
754+
finally:
755+
# Clean up additional subnets
756+
await ec2_client.delete_subnet(SubnetId=subnet2_id)
757+
await ec2_client.delete_subnet(SubnetId=subnet3_id)
758+
759+
760+
async def test_launch_instances_insufficient_capacity_fallback(
761+
simcore_ec2_api: SimcoreEC2API,
762+
ec2_client: EC2Client,
763+
fake_ec2_instance_type: EC2InstanceType,
764+
faker: Faker,
765+
aws_subnet_id: str,
766+
aws_security_group_id: str,
767+
aws_ami_id: str,
768+
):
769+
"""Test that launch_instances falls back to next subnet when InsufficientInstanceCapacity occurs."""
770+
await _assert_no_instances_in_ec2(ec2_client)
771+
772+
# Create additional valid subnets for testing
773+
vpc_id = (await ec2_client.describe_subnets(SubnetIds=[aws_subnet_id]))["Subnets"][
774+
0
775+
]["VpcId"]
776+
777+
# Create second subnet
778+
subnet2 = await ec2_client.create_subnet(CidrBlock="10.0.2.0/24", VpcId=vpc_id)
779+
subnet2_id = subnet2["Subnet"]["SubnetId"]
780+
781+
try:
782+
# Create a config with multiple valid subnet IDs
783+
ec2_instance_config = EC2InstanceConfig(
784+
type=fake_ec2_instance_type,
785+
tags=faker.pydict(allowed_types=(str,)),
786+
startup_script=faker.pystr(),
787+
ami_id=aws_ami_id,
788+
key_name=faker.pystr(),
789+
security_group_ids=[aws_security_group_id],
790+
subnet_ids=[aws_subnet_id, subnet2_id],
791+
iam_instance_profile="",
792+
)
793+
794+
# Mock the EC2 client to simulate InsufficientInstanceCapacity on first subnet
795+
original_run_instances = simcore_ec2_api.client.run_instances
796+
call_count = 0
797+
798+
async def mock_run_instances(*args, **kwargs):
799+
nonlocal call_count
800+
call_count += 1
801+
if call_count == 1:
802+
# First call (first subnet) - simulate insufficient capacity
803+
error_response = {
804+
"Error": {
805+
"Code": "InsufficientInstanceCapacity",
806+
"Message": "Insufficient capacity.",
807+
}
808+
}
809+
raise botocore.exceptions.ClientError(error_response, "RunInstances")
810+
# Second call (second subnet) - succeed normally
811+
return await original_run_instances(*args, **kwargs)
812+
813+
# Apply the mock
814+
with patch.object(
815+
simcore_ec2_api.client, "run_instances", side_effect=mock_run_instances
816+
):
817+
instances = await simcore_ec2_api.launch_instances(
818+
ec2_instance_config,
819+
min_number_of_instances=1,
820+
number_of_instances=1,
821+
)
822+
823+
# Verify that run_instances was called twice (once for each subnet)
824+
assert call_count == 2
825+
826+
# Verify that the instance was created (in the second subnet)
827+
await _assert_instances_in_ec2(
828+
ec2_client,
829+
expected_num_reservations=1,
830+
expected_num_instances=1,
831+
expected_instance_type=ec2_instance_config.type,
832+
expected_tags=ec2_instance_config.tags,
833+
expected_state="running",
834+
)
835+
836+
# Verify the instance was created in the second subnet (since first failed)
837+
instance_details = await ec2_client.describe_instances(
838+
InstanceIds=[instances[0].id]
839+
)
840+
instance = instance_details["Reservations"][0]["Instances"][0]
841+
assert instance["SubnetId"] == subnet2_id
842+
843+
finally:
844+
# Clean up additional subnet
845+
await ec2_client.delete_subnet(SubnetId=subnet2_id)
846+
847+
848+
async def test_launch_instances_all_subnets_insufficient_capacity_raises_error(
849+
simcore_ec2_api: SimcoreEC2API,
850+
ec2_client: EC2Client,
851+
fake_ec2_instance_type: EC2InstanceType,
852+
faker: Faker,
853+
aws_subnet_id: str,
854+
aws_security_group_id: str,
855+
aws_ami_id: str,
856+
):
857+
"""Test that launch_instances raises EC2InsufficientCapacityError when all subnets have insufficient capacity."""
858+
await _assert_no_instances_in_ec2(ec2_client)
859+
860+
# Create additional valid subnets for testing
861+
vpc_id = (await ec2_client.describe_subnets(SubnetIds=[aws_subnet_id]))["Subnets"][
862+
0
863+
]["VpcId"]
864+
865+
# Create second subnet
866+
subnet2 = await ec2_client.create_subnet(CidrBlock="10.0.2.0/24", VpcId=vpc_id)
867+
subnet2_id = subnet2["Subnet"]["SubnetId"]
868+
869+
try:
870+
# Create a config with multiple valid subnet IDs
871+
ec2_instance_config = EC2InstanceConfig(
872+
type=fake_ec2_instance_type,
873+
tags=faker.pydict(allowed_types=(str,)),
874+
startup_script=faker.pystr(),
875+
ami_id=aws_ami_id,
876+
key_name=faker.pystr(),
877+
security_group_ids=[aws_security_group_id],
878+
subnet_ids=[aws_subnet_id, subnet2_id],
879+
iam_instance_profile="",
880+
)
881+
882+
# Mock the EC2 client to simulate InsufficientInstanceCapacity on ALL subnets
883+
call_count = 0
884+
885+
async def mock_run_instances(*args, **kwargs):
886+
nonlocal call_count
887+
call_count += 1
888+
# Always simulate insufficient capacity
889+
error_response = {
890+
"Error": {
891+
"Code": "InsufficientInstanceCapacity",
892+
"Message": "Insufficient capacity.",
893+
}
894+
}
895+
raise botocore.exceptions.ClientError(error_response, "RunInstances")
896+
897+
# Apply the mock and expect EC2InsufficientCapacityError
898+
with (
899+
patch.object(
900+
simcore_ec2_api.client, "run_instances", side_effect=mock_run_instances
901+
),
902+
pytest.raises(EC2InsufficientCapacityError) as exc_info,
903+
):
904+
await simcore_ec2_api.launch_instances(
905+
ec2_instance_config,
906+
min_number_of_instances=1,
907+
number_of_instances=1,
908+
)
909+
910+
# Verify that run_instances was called for both subnets
911+
assert call_count == 2
912+
913+
# Verify the error contains the expected information
914+
assert exc_info.value.instance_type == fake_ec2_instance_type.name
915+
916+
# Verify no instances were created
917+
await _assert_no_instances_in_ec2(ec2_client)
918+
919+
finally:
920+
# Clean up additional subnet
921+
await ec2_client.delete_subnet(SubnetId=subnet2_id)

0 commit comments

Comments
 (0)