|
35 | 35 | from unittest import mock |
36 | 36 | from unittest.mock import MagicMock, patch |
37 | 37 |
|
38 | | -import moto |
39 | 38 | import pytest |
40 | 39 | import yaml |
41 | 40 | from click.testing import CliRunner |
42 | | -from moto import mock_ec2, mock_iam |
| 41 | +from moto import mock_aws |
43 | 42 | from testfixtures import Replacer |
44 | 43 | from testfixtures.popen import MockPopen, PopenBehaviour |
45 | 44 |
|
@@ -113,18 +112,25 @@ def configure_aws(): |
113 | 112 | os.environ["AWS_SESSION_TOKEN"] = "testing" |
114 | 113 |
|
115 | 114 | # moto (boto3 mock) only allows a hardcoded set of AMIs |
116 | | - dlami = ( |
117 | | - moto.ec2.models.ec2_backends["us-west-2"]["us-west-2"] |
118 | | - .describe_images(filters={"name": "Deep Learning AMI Ubuntu*"})[0] |
119 | | - .id |
120 | | - ) |
121 | | - aws_config.DEFAULT_AMI["us-west-2"] = dlami |
122 | | - list_instances_mock = MagicMock(return_value=boto3_list) |
123 | | - with patch( |
124 | | - "ray.autoscaler._private.aws.node_provider.list_ec2_instances", |
125 | | - list_instances_mock, |
126 | | - ): |
127 | | - yield |
| 115 | + # Use mock_aws context manager and boto3 to find the AMI |
| 116 | + import boto3 |
| 117 | + |
| 118 | + # In moto 5.x, AWS managed policies (e.g., AmazonEC2FullAccess) are not |
| 119 | + # loaded by default for performance. Enable them since the autoscaler |
| 120 | + # attaches these policies to the IAM role. |
| 121 | + with mock_aws(config={"iam": {"load_aws_managed_policies": True}}): |
| 122 | + ec2_client = boto3.client("ec2", region_name="us-west-2") |
| 123 | + images = ec2_client.describe_images( |
| 124 | + Filters=[{"Name": "name", "Values": ["Deep Learning AMI Ubuntu*"]}] |
| 125 | + )["Images"] |
| 126 | + dlami = images[0]["ImageId"] |
| 127 | + aws_config.DEFAULT_AMI["us-west-2"] = dlami |
| 128 | + list_instances_mock = MagicMock(return_value=boto3_list) |
| 129 | + with patch( |
| 130 | + "ray.autoscaler._private.aws.node_provider.list_ec2_instances", |
| 131 | + list_instances_mock, |
| 132 | + ): |
| 133 | + yield |
128 | 134 |
|
129 | 135 |
|
130 | 136 | @pytest.fixture(scope="function") |
@@ -636,8 +642,6 @@ def test_ray_start_block_and_stop(configure_lang, monkeypatch, tmp_path, cleanup |
636 | 642 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
637 | 643 | reason=("Mac builds don't provide proper locale support"), |
638 | 644 | ) |
639 | | -@mock_ec2 |
640 | | -@mock_iam |
641 | 645 | def test_ray_up( |
642 | 646 | configure_lang, _unlink_test_ssh_key, configure_aws, monkeypatch, tmp_path |
643 | 647 | ): |
@@ -677,8 +681,6 @@ def commands_mock(command, stdin): |
677 | 681 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
678 | 682 | reason=("Mac builds don't provide proper locale support"), |
679 | 683 | ) |
680 | | -@mock_ec2 |
681 | | -@mock_iam |
682 | 684 | def test_ray_up_docker( |
683 | 685 | configure_lang, _unlink_test_ssh_key, configure_aws, monkeypatch, tmp_path |
684 | 686 | ): |
@@ -720,8 +722,6 @@ def commands_mock(command, stdin): |
720 | 722 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
721 | 723 | reason=("Mac builds don't provide proper locale support"), |
722 | 724 | ) |
723 | | -@mock_ec2 |
724 | | -@mock_iam |
725 | 725 | def test_ray_up_record( |
726 | 726 | configure_lang, _unlink_test_ssh_key, configure_aws, monkeypatch, tmp_path |
727 | 727 | ): |
@@ -754,8 +754,6 @@ def commands_mock(command, stdin): |
754 | 754 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
755 | 755 | reason=("Mac builds don't provide proper locale support"), |
756 | 756 | ) |
757 | | -@mock_ec2 |
758 | | -@mock_iam |
759 | 757 | def test_ray_attach(configure_lang, configure_aws, _unlink_test_ssh_key): |
760 | 758 | def commands_mock(command, stdin): |
761 | 759 | # TODO(maximsmol): this is a hack since stdout=sys.stdout |
@@ -796,8 +794,6 @@ def commands_mock(command, stdin): |
796 | 794 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
797 | 795 | reason=("Mac builds don't provide proper locale support"), |
798 | 796 | ) |
799 | | -@mock_ec2 |
800 | | -@mock_iam |
801 | 797 | def test_ray_attach_with_ip(configure_lang, configure_aws, _unlink_test_ssh_key): |
802 | 798 | from ray.autoscaler._private.commands import get_worker_node_ips |
803 | 799 |
|
@@ -876,8 +872,6 @@ def commands_verifier(calls): |
876 | 872 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
877 | 873 | reason=("Mac builds don't provide proper locale support"), |
878 | 874 | ) |
879 | | -@mock_ec2 |
880 | | -@mock_iam |
881 | 875 | def test_ray_dashboard(configure_lang, configure_aws, _unlink_test_ssh_key): |
882 | 876 | def commands_mock(command, stdin): |
883 | 877 | # TODO(maximsmol): this is a hack since stdout=sys.stdout |
@@ -910,8 +904,6 @@ def commands_mock(command, stdin): |
910 | 904 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
911 | 905 | reason=("Mac builds don't provide proper locale support"), |
912 | 906 | ) |
913 | | -@mock_ec2 |
914 | | -@mock_iam |
915 | 907 | def test_ray_exec(configure_lang, configure_aws, _unlink_test_ssh_key): |
916 | 908 | def commands_mock(command, stdin): |
917 | 909 | # TODO(maximsmol): this is a hack since stdout=sys.stdout |
@@ -963,8 +955,6 @@ def commands_verifier(calls): |
963 | 955 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
964 | 956 | reason=("Mac builds don't provide proper locale support"), |
965 | 957 | ) |
966 | | -@mock_ec2 |
967 | | -@mock_iam |
968 | 958 | def test_ray_submit(configure_lang, configure_aws, _unlink_test_ssh_key): |
969 | 959 | def commands_mock(command, stdin): |
970 | 960 | # TODO(maximsmol): this is a hack since stdout=sys.stdout |
@@ -1355,8 +1345,6 @@ def test_ray_drain_node(monkeypatch): |
1355 | 1345 | sys.platform == "darwin" and "travis" in os.environ.get("USER", ""), |
1356 | 1346 | reason=("Mac builds don't provide proper locale support"), |
1357 | 1347 | ) |
1358 | | -@mock_ec2 |
1359 | | -@mock_iam |
1360 | 1348 | def test_ray_cluster_dump(configure_lang, configure_aws, _unlink_test_ssh_key): |
1361 | 1349 | def commands_mock(command, stdin): |
1362 | 1350 | print("This is a test!") |
|
0 commit comments