Merge "New feature to lease a device from device pool."

9 files changed, 225 insertions, 6 deletions

diff --git a/create/avd_spec.py b/create/avd_spec.py
index 872c5b08..a03631d4 100644
--- a/create/avd_spec.py
+++ b/create/avd_spec.py
@@ -119,6 +119,7 @@ class AVDSpec():
         self._num_of_instances = None
         self._num_avds_per_instance = None
         self._no_pull_log = None
+        self._oxygen = None
         self._remote_image = None
         self._system_build_info = None
         self._kernel_build_info = None
@@ -318,6 +319,7 @@ class AVDSpec():
         self._num_of_instances = args.num
         self._num_avds_per_instance = args.num_avds_per_instance
         self._no_pull_log = args.no_pull_log
+        self._oxygen = args.oxygen
         self._serial_log_file = args.serial_log_file
         self._emulator_build_id = args.emulator_build_id
         self._gpu = args.gpu
@@ -925,3 +927,8 @@ class AVDSpec():
     def gce_metadata(self):
         """Return gce_metadata."""
         return self._gce_metadata
+
+    @property
+    def oxygen(self):
+        """Return oxygen."""
+        return self._oxygen
diff --git a/create/avd_spec_test.py b/create/avd_spec_test.py
index f337aab0..a42d52fa 100644
--- a/create/avd_spec_test.py
+++ b/create/avd_spec_test.py
@@ -451,6 +451,10 @@ class AvdSpecTest(driver_test_lib.BaseDriverTest):
         self.AvdSpec._ProcessMiscArgs(self.args)
         self.assertEqual(self.AvdSpec._instance_type, constants.INSTANCE_TYPE_HOST)
 
+        self.args.oxygen = True
+        self.AvdSpec._ProcessMiscArgs(self.args)
+        self.assertTrue(self.AvdSpec._oxygen)
+
         # Test avd_spec.autoconnect
         self.args.autoconnect = False
         self.AvdSpec._ProcessMiscArgs(self.args)
diff --git a/create/create_args.py b/create/create_args.py
index c85dc51d..bdc222ed 100644
--- a/create/create_args.py
+++ b/create/create_args.py
@@ -247,6 +247,12 @@ def AddCommonCreateArgs(parser):
         default=1,
         help=argparse.SUPPRESS)
     parser.add_argument(
+        "--oxygen",
+        action="store_true",
+        dest="oxygen",
+        required=False,
+        help=argparse.SUPPRESS)
+    parser.add_argument(
         "--zone",
         type=str,
         dest="zone",
diff --git a/create/remote_image_remote_instance.py b/create/remote_image_remote_instance.py
index 80440415..b56c411a 100644
--- a/create/remote_image_remote_instance.py
+++ b/create/remote_image_remote_instance.py
@@ -18,11 +18,23 @@ r"""RemoteImageRemoteInstance class.
 Create class that is responsible for creating a remote instance AVD with a
 remote image.
 """
+
+import logging
+import time
+
 from acloud.create import base_avd_create
+from acloud.internal import constants
+from acloud.internal.lib import engprod_client
 from acloud.internal.lib import utils
 from acloud.public.actions import common_operations
 from acloud.public.actions import remote_instance_cf_device_factory
-from acloud.internal import constants
+from acloud.public import report
+
+
+logger = logging.getLogger(__name__)
+_DEVICE = "device"
+_DEVICE_KEY_MAPPING = {"serverUrl": "ip", "sessionId": "instance_name"}
+_LAUNCH_CVD_TIME = "launch_cvd_time"
 
 
 class RemoteImageRemoteInstance(base_avd_create.BaseAVDCreate):
@@ -40,9 +52,11 @@ class RemoteImageRemoteInstance(base_avd_create.BaseAVDCreate):
         Returns:
             A Report instance.
         """
+        if avd_spec.oxygen:
+            return self._LeaseOxygenAVD(avd_spec)
         device_factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
             avd_spec)
-        report = common_operations.CreateDevices(
+        create_report = common_operations.CreateDevices(
             "create_cf", avd_spec.cfg, device_factory, avd_spec.num,
             report_internal_ip=avd_spec.report_internal_ip,
             autoconnect=avd_spec.autoconnect,
@@ -54,8 +68,55 @@ class RemoteImageRemoteInstance(base_avd_create.BaseAVDCreate):
             client_adb_port=avd_spec.client_adb_port)
         # Launch vnc client if we're auto-connecting.
         if avd_spec.connect_vnc:
-            utils.LaunchVNCFromReport(report, avd_spec, no_prompts)
+            utils.LaunchVNCFromReport(create_report, avd_spec, no_prompts)
         if avd_spec.connect_webrtc:
-            utils.LaunchBrowserFromReport(report)
+            utils.LaunchBrowserFromReport(create_report)
+
+        return create_report
+
+    def _LeaseOxygenAVD(self, avd_spec):
+        """Lease the AVD from the AVD pool.
 
-        return report
+        Args:
+            avd_spec: AVDSpec object that tells us what we're going to create.
+
+        Returns:
+            A Report instance.
+        """
+        timestart = time.time()
+        response = engprod_client.EngProdClient.LeaseDevice(
+            avd_spec.remote_image[constants.BUILD_TARGET],
+            avd_spec.remote_image[constants.BUILD_ID],
+            avd_spec.cfg.api_key,
+            avd_spec.cfg.api_url)
+        execution_time = round(time.time() - timestart, 2)
+        reporter = report.Report(command="create_cf")
+        if _DEVICE in response:
+            reporter.SetStatus(report.Status.SUCCESS)
+            device_data = response[_DEVICE]
+            device_data[_LAUNCH_CVD_TIME] = execution_time
+            self._ReplaceDeviceDataKeys(device_data)
+            reporter.UpdateData(response)
+        else:
+            reporter.SetStatus(report.Status.FAIL)
+            reporter.AddError(response.get("errorMessage"))
+
+        return reporter
+
+    @staticmethod
+    def _ReplaceDeviceDataKeys(device_data):
+        """Replace keys of device data from oxygen response.
+
+        To keep the device data using the same keys in Acloud report. Before
+        writing data to report, it needs to update the keys.
+
+        Values:
+            device_data: Dict of device data. e.g. {'sessionId': 'b01ead68',
+                                                    'serverUrl': '10.1.1.1'}
+        """
+        for key, val in _DEVICE_KEY_MAPPING.items():
+            if key in device_data:
+                device_data[val] = device_data[key]
+                del device_data[key]
+            else:
+                logger.debug("There is no '%s' data in response.", key)
diff --git a/create/remote_image_remote_instance_test.py b/create/remote_image_remote_instance_test.py
new file mode 100644
index 00000000..ba8d8107
--- /dev/null
+++ b/create/remote_image_remote_instance_test.py
@@ -0,0 +1,86 @@
+# Copyright 2021 - The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for RemoteImageRemoteInstance."""
+
+import unittest
+
+from unittest import mock
+
+from acloud.create import remote_image_remote_instance
+from acloud.internal import constants
+from acloud.internal.lib import driver_test_lib
+from acloud.internal.lib import engprod_client
+from acloud.public import report
+from acloud.public.actions import common_operations
+from acloud.public.actions import remote_instance_cf_device_factory
+
+
+class RemoteImageRemoteInstanceTest(driver_test_lib.BaseDriverTest):
+    """Test RemoteImageRemoteInstance method."""
+
+    def setUp(self):
+        """Initialize new RemoteImageRemoteInstance."""
+        super().setUp()
+        self.remote_image_remote_instance = remote_image_remote_instance.RemoteImageRemoteInstance()
+
+    # pylint: disable=protected-access
+    @mock.patch.object(remote_image_remote_instance.RemoteImageRemoteInstance,
+                       "_LeaseOxygenAVD")
+    @mock.patch.object(common_operations, "CreateDevices")
+    @mock.patch.object(remote_instance_cf_device_factory,
+                       "RemoteInstanceDeviceFactory")
+    def testCreateAVD(self, mock_factory, mock_create_device, mock_lease):
+        """test CreateAVD."""
+        avd_spec = mock.Mock()
+        avd_spec.oxygen = False
+        self.remote_image_remote_instance._CreateAVD(
+            avd_spec, no_prompts=True)
+        mock_factory.assert_called_once()
+        mock_create_device.assert_called_once()
+
+        avd_spec.oxygen = True
+        self.remote_image_remote_instance._CreateAVD(
+            avd_spec, no_prompts=True)
+        mock_lease.assert_called_once()
+
+    def testLeaseOxygenAVD(self):
+        """test LeaseOxygenAVD."""
+        avd_spec = mock.Mock()
+        avd_spec.oxygen = True
+        avd_spec.remote_image = {constants.BUILD_TARGET: "fake_target",
+                                 constants.BUILD_ID: "fake_id"}
+        response_success = {"device": {"sessionId": "fake_device",
+                                       "serverUrl": "10.1.1.1"}}
+        response_fail = {"errorMessage": "Lease device fail."}
+        self.Patch(engprod_client.EngProdClient, "LeaseDevice",
+                   side_effect=[response_success, response_fail])
+        expected_status = report.Status.SUCCESS
+        reporter = self.remote_image_remote_instance._LeaseOxygenAVD(avd_spec)
+        self.assertEqual(reporter.status, expected_status)
+
+        expected_status = report.Status.FAIL
+        reporter = self.remote_image_remote_instance._LeaseOxygenAVD(avd_spec)
+        self.assertEqual(reporter.status, expected_status)
+
+
+    def testReplaceDeviceDataKeys(self):
+        """test ReplaceDeviceDataKeys."""
+        device_data = {"sessionId": "fake_device", "serverUrl": "10.1.1.1"}
+        expected_result = {"instance_name": "fake_device", "ip": "10.1.1.1"}
+        self.remote_image_remote_instance._ReplaceDeviceDataKeys(device_data)
+        self.assertEqual(device_data, expected_result)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/internal/lib/base_cloud_client.py b/internal/lib/base_cloud_client.py
index b4167c5c..6e4400c5 100755
--- a/internal/lib/base_cloud_client.py
+++ b/internal/lib/base_cloud_client.py
@@ -37,7 +37,7 @@ from acloud.internal.lib import utils
 logger = logging.getLogger(__name__)
 
 
-class BaseCloudApiClient(object):
+class BaseCloudApiClient():
     """A class that does basic setup for a cloud API."""
 
     # To be overriden by subclasses.
diff --git a/internal/lib/engprod_client.py b/internal/lib/engprod_client.py
new file mode 100644
index 00000000..26043543
--- /dev/null
+++ b/internal/lib/engprod_client.py
@@ -0,0 +1,47 @@
+# Copyright 2021 - The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A client that talks to EngProd APIs."""
+
+import json
+import subprocess
+
+from urllib.parse import urljoin
+
+
+class EngProdClient():
+    """Client that manages EngProd api."""
+
+    @staticmethod
+    def LeaseDevice(build_target, build_id, api_key, api_url):
+        """Lease one cuttlefish device.
+
+        Args:
+            build_target: Target name, e.g. "aosp_cf_x86_phone-userdebug"
+            build_id: Build ID, a string, e.g. "2263051", "P2804227"
+            api_key: String of api key.
+            api_url: String of api url.
+
+        Returns:
+            The response of curl command.
+        """
+        request_data = "{\"target\": \"%s\", \"build_id\": \"%s\"}" % (
+            build_target, build_id)
+        lease_url = urljoin(api_url, "lease?key=%s" % api_key)
+        response = subprocess.check_output([
+            "curl", "--request", "POST", lease_url, "-H",
+            "Accept: application/json", "-H", "Content-Type: application/json",
+            "-d", request_data
+        ])
+        return json.loads(response)
diff --git a/internal/proto/user_config.proto b/internal/proto/user_config.proto
index c83f15cd..0dec7717 100755
--- a/internal/proto/user_config.proto
+++ b/internal/proto/user_config.proto
@@ -112,4 +112,10 @@ message UserConfig {
 
   // [CHEEPS only] The name of the L1 betty image (used with Cheeps controller)
   optional string betty_image = 31;
+
+  // [Oxygen only] The OAuth Credentials of API key.
+  optional string api_key = 32;
+
+  // [Oxygen only] The API service url.
+  optional string api_url = 33;
 }
diff --git a/public/config.py b/public/config.py
index 5b27a437..51517a56 100755
--- a/public/config.py
+++ b/public/config.py
@@ -232,6 +232,8 @@ class AcloudConfig():
         self.hw_property = usr_cfg.hw_property
 
         self.launch_args = usr_cfg.launch_args
+        self.api_key = usr_cfg.api_key
+        self.api_url = usr_cfg.api_url
         self.instance_name_pattern = (
             usr_cfg.instance_name_pattern or
             internal_cfg.default_usr_cfg.instance_name_pattern)