From 51d277f2e4c33cf0b9fbd89f631c0362b1c57cef Mon Sep 17 00:00:00 2001
From: Richard Liaw <rliaw@berkeley.edu>
Date: Mon, 16 Nov 2020 18:44:15 -0800
Subject: [PATCH] [tests] fix mock for test_cli (#12055)

Signed-off-by: Richard Liaw <rliaw@berkeley.edu>
---
 python/ray/tests/test_cli.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/python/ray/tests/test_cli.py b/python/ray/tests/test_cli.py
index e4427a00a..696a0271a 100644
--- a/python/ray/tests/test_cli.py
+++ b/python/ray/tests/test_cli.py
@@ -29,6 +29,7 @@ import pytest
 
 import moto
 from moto import mock_ec2, mock_iam
+from unittest.mock import MagicMock, patch
 from click.testing import CliRunner
 
 from testfixtures import Replacer
@@ -37,6 +38,34 @@ from testfixtures.popen import MockPopen, PopenBehaviour
 import ray.autoscaler._private.aws.config as aws_config
 import ray.scripts.scripts as scripts
 
+boto3_list = [{
+    "InstanceType": "t1.micro",
+    "VCpuInfo": {
+        "DefaultVCpus": 1
+    }
+}, {
+    "InstanceType": "m4.xlarge",
+    "VCpuInfo": {
+        "DefaultVCpus": 4
+    }
+}, {
+    "InstanceType": "m4.4xlarge",
+    "VCpuInfo": {
+        "DefaultVCpus": 16
+    }
+}, {
+    "InstanceType": "p3.8xlarge",
+    "VCpuInfo": {
+        "DefaultVCpus": 32
+    },
+    "GpuInfo": {
+        "Gpus": [{
+            "Name": "V100",
+            "Count": 4
+        }]
+    }
+}]
+
 
 @pytest.fixture
 def configure_lang():
@@ -60,6 +89,10 @@ def configure_aws():
     dlami = moto.ec2.ec2_backends["us-west-2"].describe_images(
         filters={"name": "Deep Learning AMI Ubuntu*"})[0].id
     aws_config.DEFAULT_AMI["us-west-2"] = dlami
+    list_instances_mock = MagicMock(return_value=boto3_list)
+    with patch("ray.autoscaler._private.aws.node_provider.list_ec2_instances",
+               list_instances_mock):
+        yield
 
 
 @pytest.fixture(scope="function")