add internal/lib/android_build_client.py and

internal/lib/android_build_client_test.py

2 files changed, 260 insertions, 0 deletions

diff --git a/internal/lib/android_build_client.py b/internal/lib/android_build_client.py
new file mode 100644
index 00000000..b459f621
--- /dev/null
+++ b/internal/lib/android_build_client.py
@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 - 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 Android Build APIs."""
+
+import io
+import logging
+import os
+
+import apiclient
+
+from acloud.internal.lib import base_cloud_client
+from acloud.public import errors
+
+logger = logging.getLogger(__name__)
+
+
+class AndroidBuildClient(base_cloud_client.BaseCloudApiClient):
+    """Client that manages Android Build."""
+
+    # API settings, used by BaseCloudApiClient.
+    API_NAME = "androidbuildinternal"
+    API_VERSION = "v2beta1"
+    SCOPE = "https://www.googleapis.com/auth/androidbuild.internal"
+
+    # other variables.
+    DEFAULT_RESOURCE_ID = "0"
+    # TODO(fdeng): We should use "latest".
+    DEFAULT_ATTEMPT_ID = "0"
+    DEFAULT_CHUNK_SIZE = 20 * 1024 * 1024
+    NO_ACCESS_ERROR_PATTERN = "does not have storage.objects.create access"
+
+    # Message constant
+    COPY_TO_MSG = ("build artifact (target: %s, build_id: %s, "
+                   "artifact: %s, attempt_id: %s) to "
+                   "google storage (bucket: %s, path: %s)")
+
+    def DownloadArtifact(self,
+                         build_target,
+                         build_id,
+                         resource_id,
+                         local_dest,
+                         attempt_id=None):
+        """Get Android build attempt information.
+
+        Args:
+            build_target: Target name, e.g. "gce_x86-userdebug"
+            build_id: Build id, a string, e.g. "2263051", "P2804227"
+            resource_id: Id of the resource, e.g "avd-system.tar.gz".
+            local_dest: A local path where the artifact should be stored.
+                        e.g. "/tmp/avd-system.tar.gz"
+            attempt_id: String, attempt id, will default to DEFAULT_ATTEMPT_ID.
+        """
+        attempt_id = attempt_id or self.DEFAULT_ATTEMPT_ID
+        api = self.service.buildartifact().get_media(
+            buildId=build_id,
+            target=build_target,
+            attemptId=attempt_id,
+            resourceId=resource_id)
+        logger.info("Downloading artifact: target: %s, build_id: %s, "
+                    "resource_id: %s, dest: %s", build_target, build_id,
+                    resource_id, local_dest)
+        try:
+            with io.FileIO(local_dest, mode="wb") as fh:
+                downloader = apiclient.http.MediaIoBaseDownload(
+                    fh, api, chunksize=self.DEFAULT_CHUNK_SIZE)
+                done = False
+                while not done:
+                    _, done = downloader.next_chunk()
+            logger.info("Downloaded artifact: %s", local_dest)
+        except OSError as e:
+            logger.error("Downloading artifact failed: %s", str(e))
+            raise errors.DriverError(str(e))
+
+    def CopyTo(self,
+               build_target,
+               build_id,
+               artifact_name,
+               destination_bucket,
+               destination_path,
+               attempt_id=None):
+        """Copy an Android Build artifact to a storage bucket.
+
+        Args:
+            build_target: Target name, e.g. "gce_x86-userdebug"
+            build_id: Build id, a string, e.g. "2263051", "P2804227"
+            artifact_name: Name of the artifact, e.g "avd-system.tar.gz".
+            destination_bucket: String, a google storage bucket name.
+            destination_path: String, "path/inside/bucket"
+            attempt_id: String, attempt id, will default to DEFAULT_ATTEMPT_ID.
+        """
+        attempt_id = attempt_id or self.DEFAULT_ATTEMPT_ID
+        copy_msg = "Copying %s" % self.COPY_TO_MSG
+        logger.info(copy_msg, build_target, build_id, artifact_name,
+                    attempt_id, destination_bucket, destination_path)
+        api = self.service.buildartifact().copyTo(
+            buildId=build_id,
+            target=build_target,
+            attemptId=attempt_id,
+            artifactName=artifact_name,
+            destinationBucket=destination_bucket,
+            destinationPath=destination_path)
+        try:
+            self.Execute(api)
+            finish_msg = "Finished copying %s" % self.COPY_TO_MSG
+            logger.info(finish_msg, build_target, build_id, artifact_name,
+                        attempt_id, destination_bucket, destination_path)
+        except errors.HttpError as e:
+            if e.code == 503:
+                if self.NO_ACCESS_ERROR_PATTERN in str(e):
+                    error_msg = "Please grant android build team's service account "
+                    error_msg += "write access to bucket %s. Original error: %s"
+                    error_msg %= (destination_bucket, str(e))
+                    raise errors.HttpError(e.code, message=error_msg)
+            raise
diff --git a/internal/lib/android_build_client_test.py b/internal/lib/android_build_client_test.py
new file mode 100644
index 00000000..db097bc7
--- /dev/null
+++ b/internal/lib/android_build_client_test.py
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 - 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 acloud.internal.lib.android_build_client."""
+
+import io
+import time
+
+import apiclient
+import mock
+
+import unittest
+from acloud.internal.lib import android_build_client
+from acloud.internal.lib import driver_test_lib
+from acloud.public import errors
+
+
+class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
+    """Test AndroidBuildClient."""
+
+    BUILD_TARGET = "fake_target"
+    BUILD_ID = 12345
+    RESOURCE_ID = "avd-system.tar.gz"
+    LOCAL_DEST = "/fake/local/path"
+    DESTINATION_BUCKET = "fake_bucket"
+
+    def setUp(self):
+        """Set up test."""
+        super(AndroidBuildClientTest, self).setUp()
+        self.Patch(android_build_client.AndroidBuildClient,
+                   "InitResourceHandle")
+        self.client = android_build_client.AndroidBuildClient(mock.MagicMock())
+        self.client._service = mock.MagicMock()
+
+    def testDownloadArtifact(self):
+        """Test DownloadArtifact."""
+        # Create mocks.
+        mock_file = mock.MagicMock()
+        mock_file_io = mock.MagicMock()
+        mock_file_io.__enter__.return_value = mock_file
+        mock_downloader = mock.MagicMock()
+        mock_downloader.next_chunk = mock.MagicMock(
+            side_effect=[(mock.MagicMock(), False), (mock.MagicMock(), True)])
+        mock_api = mock.MagicMock()
+        self.Patch(io, "FileIO", return_value=mock_file_io)
+        self.Patch(
+            apiclient.http,
+            "MediaIoBaseDownload",
+            return_value=mock_downloader)
+        mock_resource = mock.MagicMock()
+        self.client._service.buildartifact = mock.MagicMock(
+            return_value=mock_resource)
+        mock_resource.get_media = mock.MagicMock(return_value=mock_api)
+        # Make the call to the api
+        self.client.DownloadArtifact(self.BUILD_TARGET, self.BUILD_ID,
+                                     self.RESOURCE_ID, self.LOCAL_DEST)
+        # Verify
+        mock_resource.get_media.assert_called_with(
+            buildId=self.BUILD_ID,
+            target=self.BUILD_TARGET,
+            attemptId="0",
+            resourceId=self.RESOURCE_ID)
+        io.FileIO.assert_called_with(self.LOCAL_DEST, mode="wb")
+        mock_call = mock.call(
+            mock_file,
+            mock_api,
+            chunksize=android_build_client.AndroidBuildClient.
+            DEFAULT_CHUNK_SIZE)
+        apiclient.http.MediaIoBaseDownload.assert_has_calls([mock_call])
+        self.assertEqual(mock_downloader.next_chunk.call_count, 2)
+
+    def testDownloadArtifactOSError(self):
+        """Test DownloadArtifact when OSError is raised."""
+        self.Patch(io, "FileIO", side_effect=OSError("fake OSError"))
+        self.assertRaises(errors.DriverError, self.client.DownloadArtifact,
+                          self.BUILD_TARGET, self.BUILD_ID, self.RESOURCE_ID,
+                          self.LOCAL_DEST)
+
+    def testCopyTo(self):
+        """Test CopyTo."""
+        mock_resource = mock.MagicMock()
+        self.client._service.buildartifact = mock.MagicMock(
+            return_value=mock_resource)
+        self.client.CopyTo(
+            build_target=self.BUILD_TARGET,
+            build_id=self.BUILD_ID,
+            artifact_name=self.RESOURCE_ID,
+            destination_bucket=self.DESTINATION_BUCKET,
+            destination_path=self.RESOURCE_ID)
+        mock_resource.copyTo.assert_called_once_with(
+            buildId=self.BUILD_ID,
+            target=self.BUILD_TARGET,
+            attemptId=self.client.DEFAULT_ATTEMPT_ID,
+            artifactName=self.RESOURCE_ID,
+            destinationBucket=self.DESTINATION_BUCKET,
+            destinationPath=self.RESOURCE_ID)
+
+    def testCopyToWithRetry(self):
+        """Test CopyTo with retry."""
+        self.Patch(time, "sleep")
+        mock_resource = mock.MagicMock()
+        mock_api_request = mock.MagicMock()
+        mock_resource.copyTo.return_value = mock_api_request
+        self.client._service.buildartifact.return_value = mock_resource
+        mock_api_request.execute.side_effect = errors.HttpError(503,
+                                                                "fake error")
+        self.assertRaises(
+            errors.HttpError,
+            self.client.CopyTo,
+            build_id=self.BUILD_ID,
+            build_target=self.BUILD_TARGET,
+            artifact_name=self.RESOURCE_ID,
+            destination_bucket=self.DESTINATION_BUCKET,
+            destination_path=self.RESOURCE_ID)
+        self.assertEqual(mock_api_request.execute.call_count, 6)
+
+
+if __name__ == "__main__":
+    unittest.main()