1# Copyright 2019 - The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Tests for remote_instance_cf_device_factory."""
15
16import glob
17import os
18import shutil
19import tempfile
20import unittest
21import uuid
22
23import mock
24
25from acloud.create import avd_spec
26from acloud.internal import constants
27from acloud.create import create_common
28from acloud.internal.lib import android_build_client
29from acloud.internal.lib import auth
30from acloud.internal.lib import cvd_compute_client_multi_stage
31from acloud.internal.lib import driver_test_lib
32from acloud.internal.lib import ssh
33from acloud.internal.lib import utils
34from acloud.list import list as list_instances
35from acloud.public.actions import remote_instance_cf_device_factory
36
37
38class RemoteInstanceDeviceFactoryTest(driver_test_lib.BaseDriverTest):
39    """Test RemoteInstanceDeviceFactory method."""
40
41    def setUp(self):
42        """Set up the test."""
43        super(RemoteInstanceDeviceFactoryTest, self).setUp()
44        self.Patch(auth, "CreateCredentials", return_value=mock.MagicMock())
45        self.Patch(android_build_client.AndroidBuildClient, "InitResourceHandle")
46        self.Patch(cvd_compute_client_multi_stage.CvdComputeClient, "InitResourceHandle")
47        self.Patch(list_instances, "GetInstancesFromInstanceNames", return_value=mock.MagicMock())
48        self.Patch(list_instances, "ChooseOneRemoteInstance", return_value=mock.MagicMock())
49        self.Patch(utils, "GetBuildEnvironmentVariable",
50                   return_value="test_environ")
51        self.Patch(glob, "glob", return_vale=["fake.img"])
52
53    # pylint: disable=protected-access
54    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
55                       "_FetchBuild")
56    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
57                       "_UploadArtifacts")
58    def testProcessArtifacts(self, mock_upload, mock_download):
59        """test ProcessArtifacts."""
60        # Test image source type is local.
61        args = mock.MagicMock()
62        args.config_file = ""
63        args.avd_type = constants.TYPE_CF
64        args.flavor = "phone"
65        args.local_image = None
66        avd_spec_local_img = avd_spec.AVDSpec(args)
67        fake_image_name = "/fake/aosp_cf_x86_phone-img-eng.username.zip"
68        fake_host_package_name = "/fake/host_package.tar.gz"
69        factory_local_img = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
70            avd_spec_local_img,
71            fake_image_name,
72            fake_host_package_name)
73        factory_local_img._ProcessArtifacts(constants.IMAGE_SRC_LOCAL)
74        self.assertEqual(mock_upload.call_count, 1)
75
76        # Test image source type is remote.
77        args.local_image = ""
78        args.build_id = "1234"
79        args.branch = "fake_branch"
80        args.build_target = "fake_target"
81        args.system_build_id = "2345"
82        args.system_branch = "sys_branch"
83        args.system_build_target = "sys_target"
84        args.kernel_build_id = "3456"
85        args.kernel_branch = "kernel_branch"
86        args.kernel_build_target = "kernel_target"
87        avd_spec_remote_img = avd_spec.AVDSpec(args)
88        self.Patch(cvd_compute_client_multi_stage.CvdComputeClient, "UpdateFetchCvd")
89        factory_remote_img = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
90            avd_spec_remote_img)
91        factory_remote_img._ProcessArtifacts(constants.IMAGE_SRC_REMOTE)
92        self.assertEqual(mock_download.call_count, 1)
93
94    # pylint: disable=protected-access
95    @mock.patch.dict(os.environ, {constants.ENV_BUILD_TARGET:'fake-target'})
96    def testCreateGceInstanceNameMultiStage(self):
97        """test create gce instance."""
98        # Mock uuid
99        args = mock.MagicMock()
100        args.config_file = ""
101        args.avd_type = constants.TYPE_CF
102        args.flavor = "phone"
103        args.local_image = None
104        args.adb_port = None
105        fake_avd_spec = avd_spec.AVDSpec(args)
106        fake_avd_spec.cfg.enable_multi_stage = True
107        fake_avd_spec._instance_name_to_reuse = None
108
109        fake_uuid = mock.MagicMock(hex="1234")
110        self.Patch(uuid, "uuid4", return_value=fake_uuid)
111        self.Patch(cvd_compute_client_multi_stage.CvdComputeClient, "CreateInstance")
112        fake_host_package_name = "/fake/host_package.tar.gz"
113        fake_image_name = "/fake/aosp_cf_x86_phone-img-eng.username.zip"
114
115        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
116            fake_avd_spec,
117            fake_image_name,
118            fake_host_package_name)
119        self.assertEqual(factory._CreateGceInstance(), "ins-1234-userbuild-aosp-cf-x86-phone")
120
121        # Can't get target name from zip file name.
122        fake_image_name = "/fake/aosp_cf_x86_phone.username.zip"
123        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
124            fake_avd_spec,
125            fake_image_name,
126            fake_host_package_name)
127        self.assertEqual(factory._CreateGceInstance(), "ins-1234-userbuild-fake-target")
128
129        # No image zip path, it uses local build images.
130        fake_image_name = ""
131        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
132            fake_avd_spec,
133            fake_image_name,
134            fake_host_package_name)
135        self.assertEqual(factory._CreateGceInstance(), "ins-1234-userbuild-fake-target")
136
137    @mock.patch.dict(os.environ, {constants.ENV_BUILD_TARGET:'fake-target'})
138    def testRemoteHostInstanceName(self):
139        """Test Remote host instance name."""
140        args = mock.MagicMock()
141        args.config_file = ""
142        args.avd_type = constants.TYPE_CF
143        args.flavor = "phone"
144        args.remote_host = "1.1.1.1"
145        args.local_image = None
146        args.adb_port = None
147        fake_avd_spec = avd_spec.AVDSpec(args)
148        fake_avd_spec.cfg.enable_multi_stage = True
149        fake_avd_spec._instance_name_to_reuse = None
150        fake_uuid = mock.MagicMock(hex="1234")
151        self.Patch(uuid, "uuid4", return_value=fake_uuid)
152        self.Patch(cvd_compute_client_multi_stage.CvdComputeClient, "CreateInstance")
153        self.Patch(cvd_compute_client_multi_stage.CvdComputeClient, "InitRemoteHost")
154        fake_host_package_name = "/fake/host_package.tar.gz"
155
156        fake_image_name = "/fake/aosp_cf_x86_phone-img-eng.username.zip"
157        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
158            fake_avd_spec,
159            fake_image_name,
160            fake_host_package_name)
161        self.assertEqual(factory._InitRemotehost(), "host-1.1.1.1-userbuild-aosp_cf_x86_phone")
162
163        # No image zip path, it uses local build images.
164        fake_image_name = ""
165        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
166            fake_avd_spec,
167            fake_image_name,
168            fake_host_package_name)
169        self.assertEqual(factory._InitRemotehost(), "host-1.1.1.1-userbuild-fake-target")
170
171    def testReuseInstanceNameMultiStage(self):
172        """Test reuse instance name."""
173        args = mock.MagicMock()
174        args.config_file = ""
175        args.avd_type = constants.TYPE_CF
176        args.flavor = "phone"
177        args.local_image = None
178        args.adb_port = None
179        fake_avd_spec = avd_spec.AVDSpec(args)
180        fake_avd_spec.cfg.enable_multi_stage = True
181        fake_avd_spec._instance_name_to_reuse = "fake-1234-userbuild-fake-target"
182        fake_uuid = mock.MagicMock(hex="1234")
183        self.Patch(uuid, "uuid4", return_value=fake_uuid)
184        self.Patch(cvd_compute_client_multi_stage.CvdComputeClient, "CreateInstance")
185        fake_host_package_name = "/fake/host_package.tar.gz"
186        fake_image_name = "/fake/aosp_cf_x86_phone-img-eng.username.zip"
187        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
188            fake_avd_spec,
189            fake_image_name,
190            fake_host_package_name)
191        self.assertEqual(factory._CreateGceInstance(), "fake-1234-userbuild-fake-target")
192
193    def testGetBuildInfoDict(self):
194        """Test GetBuildInfoDict."""
195        fake_host_package_name = "/fake/host_package.tar.gz"
196        fake_image_name = "/fake/aosp_cf_x86_phone-img-eng.username.zip"
197        args = mock.MagicMock()
198        # Test image source type is local.
199        args.config_file = ""
200        args.avd_type = constants.TYPE_CF
201        args.flavor = "phone"
202        args.local_image = "fake_local_image"
203        args.adb_port = None
204        avd_spec_local_image = avd_spec.AVDSpec(args)
205        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
206            avd_spec_local_image,
207            fake_image_name,
208            fake_host_package_name)
209        self.assertEqual(factory.GetBuildInfoDict(), None)
210
211        # Test image source type is remote.
212        args.local_image = ""
213        args.build_id = "123"
214        args.branch = "fake_branch"
215        args.build_target = "fake_target"
216        args.system_build_id = "234"
217        args.system_branch = "sys_branch"
218        args.system_build_target = "sys_target"
219        args.kernel_build_id = "345"
220        args.kernel_branch = "kernel_branch"
221        args.kernel_build_target = "kernel_target"
222        avd_spec_remote_image = avd_spec.AVDSpec(args)
223        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
224            avd_spec_remote_image,
225            fake_image_name,
226            fake_host_package_name)
227        expected_build_info = {
228            "build_id": "123",
229            "build_branch": "fake_branch",
230            "build_target": "fake_target",
231            "system_build_id": "234",
232            "system_build_branch": "sys_branch",
233            "system_build_target": "sys_target",
234            "kernel_build_id": "345",
235            "kernel_build_branch": "kernel_branch",
236            "kernel_build_target": "kernel_target"
237        }
238        self.assertEqual(factory.GetBuildInfoDict(), expected_build_info)
239
240    @mock.patch.object(ssh, "ShellCmdWithRetry")
241    @mock.patch.object(ssh.Ssh, "Run")
242    def testUploadArtifacts(self, mock_ssh_run, mock_shell):
243        """Test UploadArtifacts."""
244        fake_host_package = "/fake/host_package.tar.gz"
245        fake_image = "/fake/aosp_cf_x86_phone-img-eng.username.zip"
246        fake_local_image_dir = "/fake_image"
247        fake_ip = ssh.IP(external="1.1.1.1", internal="10.1.1.1")
248        args = mock.MagicMock()
249        # Test local image extract from image zip case.
250        args.config_file = ""
251        args.avd_type = constants.TYPE_CF
252        args.flavor = "phone"
253        args.local_image = "fake_local_image"
254        args.adb_port = None
255        avd_spec_local_image = avd_spec.AVDSpec(args)
256        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
257            avd_spec_local_image,
258            fake_image,
259            fake_host_package)
260        factory._ssh = ssh.Ssh(ip=fake_ip,
261                               user=constants.GCE_USER,
262                               ssh_private_key_path="/fake/acloud_rea")
263        factory._UploadArtifacts(fake_image, fake_host_package, fake_local_image_dir)
264        expected_cmd1 = ("/usr/bin/install_zip.sh . < %s" % fake_image)
265        expected_cmd2 = ("tar -x -z -f - < %s" % fake_host_package)
266        mock_ssh_run.assert_has_calls([
267            mock.call(expected_cmd1),
268            mock.call(expected_cmd2)])
269
270        # Test local image get from local folder case.
271        fake_image = None
272        self.Patch(glob, "glob", return_value=["fake.img"])
273        factory._UploadArtifacts(fake_image, fake_host_package, fake_local_image_dir)
274        expected_cmd = (
275            "tar -cf - --lzop -S -C %s fake.img | "
276            "%s -- tar -xf - --lzop -S" %
277            (fake_local_image_dir, factory._ssh.GetBaseCmd(constants.SSH_BIN)))
278        mock_shell.assert_called_once_with(expected_cmd)
279
280    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
281                       "_InitRemotehost")
282    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
283                       "_UploadArtifacts")
284    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
285                       "_LaunchCvd")
286    def testLocalImageRemoteHost(self, mock_launchcvd, mock_upload, mock_init_remote_host):
287        """Test local image with remote host."""
288        self.Patch(
289            cvd_compute_client_multi_stage,
290            "CvdComputeClient",
291            return_value=mock.MagicMock())
292        fake_avd_spec = mock.MagicMock()
293        fake_avd_spec.instance_type = constants.INSTANCE_TYPE_HOST
294        fake_avd_spec.image_source = constants.IMAGE_SRC_LOCAL
295        fake_avd_spec._instance_name_to_reuse = None
296        fake_host_package_name = "/fake/host_package.tar.gz"
297        fake_image_name = ""
298        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
299            fake_avd_spec,
300            fake_image_name,
301            fake_host_package_name)
302        factory.CreateInstance()
303        self.assertEqual(mock_init_remote_host.call_count, 1)
304        self.assertEqual(mock_upload.call_count, 1)
305        self.assertEqual(mock_launchcvd.call_count, 1)
306
307    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
308                       "_CreateGceInstance")
309    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
310                       "_UploadArtifacts")
311    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
312                       "_LaunchCvd")
313    def testLocalImageCreateInstance(self, mock_launchcvd, mock_upload, mock_create_gce_instance):
314        """Test local image with create instance."""
315        self.Patch(
316            cvd_compute_client_multi_stage,
317            "CvdComputeClient",
318            return_value=mock.MagicMock())
319        fake_avd_spec = mock.MagicMock()
320        fake_avd_spec.instance_type = constants.INSTANCE_TYPE_REMOTE
321        fake_avd_spec.image_source = constants.IMAGE_SRC_LOCAL
322        fake_avd_spec._instance_name_to_reuse = None
323        fake_host_package_name = "/fake/host_package.tar.gz"
324        fake_image_name = ""
325        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
326            fake_avd_spec,
327            fake_image_name,
328            fake_host_package_name)
329        factory.CreateInstance()
330        self.assertEqual(mock_create_gce_instance.call_count, 1)
331        self.assertEqual(mock_upload.call_count, 1)
332        self.assertEqual(mock_launchcvd.call_count, 1)
333
334    # pylint: disable=no-member
335    @mock.patch.object(create_common, "DownloadRemoteArtifact")
336    def testDownloadArtifacts(self, mock_download):
337        """Test process remote cuttlefish image."""
338        extract_path = "/tmp/1111/"
339        fake_remote_image = {"build_target" : "aosp_cf_x86_phone-userdebug",
340                             "build_id": "1234"}
341        self.Patch(
342            cvd_compute_client_multi_stage,
343            "CvdComputeClient",
344            return_value=mock.MagicMock())
345        self.Patch(tempfile, "mkdtemp", return_value="/tmp/1111/")
346        self.Patch(shutil, "rmtree")
347        fake_avd_spec = mock.MagicMock()
348        fake_avd_spec.cfg = mock.MagicMock()
349        fake_avd_spec.remote_image = fake_remote_image
350        fake_avd_spec.image_download_dir = "/tmp"
351        self.Patch(os.path, "exists", return_value=False)
352        self.Patch(os, "makedirs")
353        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
354            fake_avd_spec)
355        factory._DownloadArtifacts(extract_path)
356        build_id = "1234"
357        build_target = "aosp_cf_x86_phone-userdebug"
358        checkfile1 = "aosp_cf_x86_phone-img-1234.zip"
359        checkfile2 = "cvd-host_package.tar.gz"
360
361        # To validate DownloadArtifact runs twice.
362        self.assertEqual(mock_download.call_count, 2)
363
364        # To validate DownloadArtifact arguments correct.
365        mock_download.assert_has_calls([
366            mock.call(fake_avd_spec.cfg, build_target, build_id, checkfile1,
367                      extract_path, decompress=True),
368            mock.call(fake_avd_spec.cfg, build_target, build_id, checkfile2,
369                      extract_path)], any_order=True)
370
371    @mock.patch.object(create_common, "DownloadRemoteArtifact")
372    @mock.patch.object(remote_instance_cf_device_factory.RemoteInstanceDeviceFactory,
373                       "_UploadArtifacts")
374    def testProcessRemoteHostArtifacts(self, mock_upload, mock_download):
375        """Test process remote host artifacts."""
376        self.Patch(
377            cvd_compute_client_multi_stage,
378            "CvdComputeClient",
379            return_value=mock.MagicMock())
380        fake_avd_spec = mock.MagicMock()
381
382        # Test process remote host artifacts with local images.
383        fake_avd_spec.instance_type = constants.INSTANCE_TYPE_HOST
384        fake_avd_spec.image_source = constants.IMAGE_SRC_LOCAL
385        fake_avd_spec._instance_name_to_reuse = None
386        fake_host_package_name = "/fake/host_package.tar.gz"
387        fake_image_name = ""
388        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
389            fake_avd_spec,
390            fake_image_name,
391            fake_host_package_name)
392        factory._ProcessRemoteHostArtifacts()
393        self.assertEqual(mock_upload.call_count, 1)
394
395        # Test process remote host artifacts with remote images.
396        fake_tmp_folder = "/tmp/1111/"
397        mock_upload.call_count = 0
398        self.Patch(tempfile, "mkdtemp", return_value=fake_tmp_folder)
399        self.Patch(shutil, "rmtree")
400        fake_avd_spec.instance_type = constants.INSTANCE_TYPE_HOST
401        fake_avd_spec.image_source = constants.IMAGE_SRC_REMOTE
402        fake_avd_spec._instance_name_to_reuse = None
403        factory = remote_instance_cf_device_factory.RemoteInstanceDeviceFactory(
404            fake_avd_spec)
405        factory._ProcessRemoteHostArtifacts()
406        self.assertEqual(mock_upload.call_count, 1)
407        self.assertEqual(mock_download.call_count, 2)
408        shutil.rmtree.assert_called_once_with(fake_tmp_folder)
409
410
411if __name__ == "__main__":
412    unittest.main()
413