Delete temporary files after acloud create with remote host. am: 0f181745c1 am: 38f4d92a3f am: abba936fd7

Change-Id: I29c2e7cb374d798c5f2b7df91079b04e0445e151

4 files changed, 84 insertions, 107 deletions

diff --git a/create/remote_image_remote_host.py b/create/remote_image_remote_host.py
index 700b1693..7ef67c4d 100644
--- a/create/remote_image_remote_host.py
+++ b/create/remote_image_remote_host.py
@@ -20,12 +20,8 @@ remote image.
 """
 
 import logging
-import os
-import shutil
-import tempfile
 
 from acloud.create import base_avd_create
-from acloud.create import create_common
 from acloud.internal import constants
 from acloud.internal.lib import utils
 from acloud.public.actions import common_operations
@@ -35,33 +31,6 @@ from acloud.public.actions import remote_instance_cf_device_factory
 logger = logging.getLogger(__name__)
 
 
-@utils.TimeExecute(function_description="Downloading Android Build artifact")
-def DownloadAndProcessArtifact(avd_spec, extract_path):
-    """Download the CF image artifacts and process them.
-
-    - Download image from the Android Build system, then decompress it.
-    - Download cvd host package from the Android Build system.
-
-    Args:
-        avd_spec: AVDSpec object that tells us what we're going to create.
-        extract_path: String, path to image folder.
-    """
-    cfg = avd_spec.cfg
-    build_id = avd_spec.remote_image[constants.BUILD_ID]
-    build_target = avd_spec.remote_image[constants.BUILD_TARGET]
-
-    logger.debug("Extract path: %s", extract_path)
-    # Image zip
-    remote_image = "%s-img-%s.zip" % (build_target.split('-')[0],
-                                      build_id)
-    create_common.DownloadRemoteArtifact(
-        cfg, build_target, build_id, remote_image, extract_path, decompress=True)
-    # Cvd host package
-    create_common.DownloadRemoteArtifact(
-        cfg, build_target, build_id, constants.CVD_HOST_PACKAGE,
-        extract_path)
-
-
 class RemoteImageRemoteHost(base_avd_create.BaseAVDCreate):
     """Create class for a remote image remote host AVD."""
 
@@ -75,13 +44,8 @@ class RemoteImageRemoteHost(base_avd_create.BaseAVDCreate):
         Returns:
             A Report instance.
         """
-        extract_path = tempfile.mkdtemp()
-        DownloadAndProcessArtifact(avd_spec, extract_path)
         device_factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
-            avd_spec=avd_spec,
-            cvd_host_package_artifact=os.path.join(
-                extract_path, constants.CVD_HOST_PACKAGE),
-            local_image_dir=extract_path)
+            avd_spec=avd_spec)
         report = common_operations.CreateDevices(
             "create_cf", avd_spec.cfg, device_factory, num=1,
             report_internal_ip=avd_spec.report_internal_ip,
@@ -93,5 +57,4 @@ class RemoteImageRemoteHost(base_avd_create.BaseAVDCreate):
         # Launch vnc client if we're auto-connecting.
         if avd_spec.connect_vnc:
             utils.LaunchVNCFromReport(report, avd_spec, no_prompts)
-        shutil.rmtree(extract_path)
         return report
diff --git a/create/remote_image_remote_host_test.py b/create/remote_image_remote_host_test.py
index abdce58e..e69de29b 100644
--- a/create/remote_image_remote_host_test.py
+++ b/create/remote_image_remote_host_test.py
@@ -1,64 +0,0 @@
-# Copyright 2018 - 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 remote_image_local_instance."""
-
-import unittest
-import os
-import mock
-
-from acloud.create import create_common
-from acloud.create import remote_image_remote_host
-from acloud.internal.lib import driver_test_lib
-
-
-# pylint: disable=invalid-name, protected-access
-class RemoteImageRemoteHostTest(driver_test_lib.BaseDriverTest):
-    """Test remote_image_local_instance methods."""
-
-    def setUp(self):
-        """Initialize remote_image_local_instance."""
-        super(RemoteImageRemoteHostTest, self).setUp()
-        self._fake_remote_image = {"build_target" : "aosp_cf_x86_phone-userdebug",
-                                   "build_id": "1234"}
-        self._extract_path = "/tmp/1111/"
-
-    @mock.patch.object(create_common, "DownloadRemoteArtifact")
-    def testDownloadAndProcessArtifact(self, mock_download):
-        """Test process remote cuttlefish image."""
-        avd_spec = mock.MagicMock()
-        avd_spec.cfg = mock.MagicMock()
-        avd_spec.remote_image = self._fake_remote_image
-        avd_spec.image_download_dir = "/tmp"
-        self.Patch(os.path, "exists", return_value=False)
-        self.Patch(os, "makedirs")
-        remote_image_remote_host.DownloadAndProcessArtifact(
-            avd_spec, self._extract_path)
-        build_id = "1234"
-        build_target = "aosp_cf_x86_phone-userdebug"
-        checkfile1 = "aosp_cf_x86_phone-img-1234.zip"
-        checkfile2 = "cvd-host_package.tar.gz"
-
-        # To validate DownloadArtifact runs twice.
-        self.assertEqual(mock_download.call_count, 2)
-
-        # To validate DownloadArtifact arguments correct.
-        mock_download.assert_has_calls([
-            mock.call(avd_spec.cfg, build_target, build_id, checkfile1,
-                      self._extract_path, decompress=True),
-            mock.call(avd_spec.cfg, build_target, build_id, checkfile2,
-                      self._extract_path)], any_order=True)
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/public/actions/remote_instance_cf_device_factory.py b/public/actions/remote_instance_cf_device_factory.py
index 04cd5027..8935aecf 100644
--- a/public/actions/remote_instance_cf_device_factory.py
+++ b/public/actions/remote_instance_cf_device_factory.py
@@ -18,8 +18,11 @@ device factory."""
 import glob
 import logging
 import os
+import shutil
+import tempfile
 
 from acloud import errors
+from acloud.create import create_common
 from acloud.internal import constants
 from acloud.internal.lib import auth
 from acloud.internal.lib import cvd_compute_client_multi_stage
@@ -48,13 +51,12 @@ class RemoteInstanceDeviceFactory(base_device_factory.BaseDeviceFactory):
         ssh: An Ssh object.
     """
     def __init__(self, avd_spec, local_image_artifact=None,
-                 cvd_host_package_artifact=None, local_image_dir=None):
+                 cvd_host_package_artifact=None):
         """Constructs a new remote instance device factory."""
         self._avd_spec = avd_spec
         self._cfg = avd_spec.cfg
         self._local_image_artifact = local_image_artifact
         self._cvd_host_package_artifact = cvd_host_package_artifact
-        self._local_image_dir = local_image_dir
         self._report_internal_ip = avd_spec.report_internal_ip
         self.credentials = auth.CreateCredentials(avd_spec.cfg)
         # Control compute_client with enable_multi_stage
@@ -142,6 +144,30 @@ class RemoteInstanceDeviceFactory(base_device_factory.BaseDeviceFactory):
             self._ssh, ip, self._avd_spec.host_user)
         return instance
 
+    @utils.TimeExecute(function_description="Downloading Android Build artifact")
+    def _DownloadArtifacts(self, extract_path):
+        """Download the CF image artifacts and process them.
+
+        - Download image from the Android Build system, then decompress it.
+        - Download cvd host package from the Android Build system.
+
+        Args:
+            extract_path: String, a path include extracted files.
+        """
+        cfg = self._avd_spec.cfg
+        build_id = self._avd_spec.remote_image[constants.BUILD_ID]
+        build_target = self._avd_spec.remote_image[constants.BUILD_TARGET]
+
+        # Image zip
+        remote_image = "%s-img-%s.zip" % (build_target.split('-')[0], build_id)
+        create_common.DownloadRemoteArtifact(
+            cfg, build_target, build_id, remote_image, extract_path, decompress=True)
+
+        # Cvd host package
+        create_common.DownloadRemoteArtifact(
+            cfg, build_target, build_id, constants.CVD_HOST_PACKAGE,
+            extract_path)
+
     def _ProcessRemoteHostArtifacts(self):
         """Process remote host artifacts.
 
@@ -151,9 +177,21 @@ class RemoteInstanceDeviceFactory(base_device_factory.BaseDeviceFactory):
           build to local and unzip it then upload to remote host, because there
           is no permission to fetch build rom on the remote host.
         """
-        self._UploadArtifacts(
-            self._local_image_artifact, self._cvd_host_package_artifact,
-            self._local_image_dir or self._avd_spec.local_image_dir)
+        if self._avd_spec.image_source == constants.IMAGE_SRC_LOCAL:
+            self._UploadArtifacts(
+                self._local_image_artifact, self._cvd_host_package_artifact,
+                self._avd_spec.local_image_dir)
+        else:
+            try:
+                artifacts_path = tempfile.mkdtemp()
+                logger.debug("Extracted path of artifacts: %s", artifacts_path)
+                self._DownloadArtifacts(artifacts_path)
+                self._UploadArtifacts(
+                    None,
+                    os.path.join(artifacts_path, constants.CVD_HOST_PACKAGE),
+                    artifacts_path)
+            finally:
+                shutil.rmtree(artifacts_path)
 
     def _ProcessArtifacts(self, image_source):
         """Process artifacts.
diff --git a/public/actions/remote_instance_cf_device_factory_test.py b/public/actions/remote_instance_cf_device_factory_test.py
index 448b93fb..ca547d77 100644
--- a/public/actions/remote_instance_cf_device_factory_test.py
+++ b/public/actions/remote_instance_cf_device_factory_test.py
@@ -15,6 +15,8 @@
 
 import glob
 import os
+import shutil
+import tempfile
 import unittest
 import uuid
 
@@ -22,6 +24,7 @@ import mock
 
 from acloud.create import avd_spec
 from acloud.internal import constants
+from acloud.create import create_common
 from acloud.internal.lib import android_build_client
 from acloud.internal.lib import auth
 from acloud.internal.lib import cvd_compute_client_multi_stage
@@ -328,6 +331,43 @@ class RemoteInstanceDeviceFactoryTest(driver_test_lib.BaseDriverTest):
         self.assertEqual(mock_upload.call_count, 1)
         self.assertEqual(mock_launchcvd.call_count, 1)
 
+    # pylint: disable=no-member
+    @mock.patch.object(create_common, "DownloadRemoteArtifact")
+    def testDownloadArtifacts(self, mock_download):
+        """Test process remote cuttlefish image."""
+        extract_path = "/tmp/1111/"
+        fake_remote_image = {"build_target" : "aosp_cf_x86_phone-userdebug",
+                             "build_id": "1234"}
+        self.Patch(
+            cvd_compute_client_multi_stage,
+            "CvdComputeClient",
+            return_value=mock.MagicMock())
+        self.Patch(tempfile, "mkdtemp", return_value="/tmp/1111/")
+        self.Patch(shutil, "rmtree")
+        fake_avd_spec = mock.MagicMock()
+        fake_avd_spec.cfg = mock.MagicMock()
+        fake_avd_spec.remote_image = fake_remote_image
+        fake_avd_spec.image_download_dir = "/tmp"
+        self.Patch(os.path, "exists", return_value=False)
+        self.Patch(os, "makedirs")
+        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
+            fake_avd_spec)
+        factory._DownloadArtifacts(extract_path)
+        build_id = "1234"
+        build_target = "aosp_cf_x86_phone-userdebug"
+        checkfile1 = "aosp_cf_x86_phone-img-1234.zip"
+        checkfile2 = "cvd-host_package.tar.gz"
+
+        # To validate DownloadArtifact runs twice.
+        self.assertEqual(mock_download.call_count, 2)
+
+        # To validate DownloadArtifact arguments correct.
+        mock_download.assert_has_calls([
+            mock.call(fake_avd_spec.cfg, build_target, build_id, checkfile1,
+                      extract_path, decompress=True),
+            mock.call(fake_avd_spec.cfg, build_target, build_id, checkfile2,
+                      extract_path)], any_order=True)
+
 
 if __name__ == "__main__":
     unittest.main()