• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2022 - 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
15"""Tests for cvd_utils."""
16
17import os
18import subprocess
19import tempfile
20import unittest
21from unittest import mock
22import zipfile
23
24from acloud import errors
25from acloud.create import create_common
26from acloud.internal import constants
27from acloud.internal.lib import cvd_utils
28from acloud.internal.lib import driver_test_lib
29
30
31# pylint: disable=too-many-public-methods
32class CvdUtilsTest(driver_test_lib.BaseDriverTest):
33    """Test the functions in cvd_utils."""
34
35    # Remote host instance name.
36    _PRODUCT_NAME = "aosp_cf_x86_64_phone"
37    _BUILD_ID = "2263051"
38    _REMOTE_HOST_IP = "192.0.2.1"
39    _REMOTE_HOST_INSTANCE_NAME_1 = (
40        "host-192.0.2.1-1-2263051-aosp_cf_x86_64_phone")
41    _REMOTE_HOST_INSTANCE_NAME_2 = (
42        "host-192.0.2.1-2-2263051-aosp_cf_x86_64_phone")
43
44    def testGetAdbPorts(self):
45        """Test GetAdbPorts."""
46        self.assertEqual([6520], cvd_utils.GetAdbPorts(None, None))
47        self.assertEqual([6520], cvd_utils.GetAdbPorts(1, 1))
48        self.assertEqual([6521, 6522], cvd_utils.GetAdbPorts(2, 2))
49
50    def testGetVncPorts(self):
51        """Test GetVncPorts."""
52        self.assertEqual([6444], cvd_utils.GetVncPorts(None, None))
53        self.assertEqual([6444], cvd_utils.GetVncPorts(1, 1))
54        self.assertEqual([6445, 6446], cvd_utils.GetVncPorts(2, 2))
55
56    def testExtractTargetFilesZip(self):
57        """Test ExtractTargetFilesZip."""
58        with tempfile.TemporaryDirectory() as temp_dir:
59            zip_path = os.path.join(temp_dir, "in.zip")
60            output_dir = os.path.join(temp_dir, "out")
61            with zipfile.ZipFile(zip_path, "w") as zip_file:
62                for entry in ["IMAGES/", "META/", "test.img",
63                              "IMAGES/system.img", "IMAGES/system.map",
64                              "IMAGES/bootloader", "IMAGES/kernel",
65                              "META/misc_info.txt"]:
66                    zip_file.writestr(entry, "")
67            cvd_utils.ExtractTargetFilesZip(zip_path, output_dir)
68
69            self.assertEqual(["IMAGES", "META"],
70                             sorted(os.listdir(output_dir)))
71            self.assertEqual(
72                ["bootloader", "kernel", "system.img"],
73                sorted(os.listdir(os.path.join(output_dir, "IMAGES"))))
74            self.assertEqual(["misc_info.txt"],
75                             os.listdir(os.path.join(output_dir, "META")))
76
77    @staticmethod
78    @mock.patch("acloud.internal.lib.cvd_utils.os.path.isdir",
79                return_value=False)
80    def testUploadImageZip(_mock_isdir):
81        """Test UploadArtifacts with image zip."""
82        mock_ssh = mock.Mock()
83        cvd_utils.UploadArtifacts(mock_ssh, "dir", "/mock/img.zip",
84                                  "/mock/cvd.tar.gz")
85        mock_ssh.Run.assert_any_call("/usr/bin/install_zip.sh dir < "
86                                     "/mock/img.zip")
87        mock_ssh.Run.assert_any_call("tar -xzf - -C dir < /mock/cvd.tar.gz")
88
89    @mock.patch("acloud.internal.lib.cvd_utils.glob")
90    @mock.patch("acloud.internal.lib.cvd_utils.os.path.isdir")
91    @mock.patch("acloud.internal.lib.cvd_utils.ssh.ShellCmdWithRetry")
92    def testUploadImageDir(self, mock_shell, mock_isdir, mock_glob):
93        """Test UploadArtifacts with image directory."""
94        mock_isdir.side_effect = lambda path: path != "/mock/cvd.tar.gz"
95        mock_ssh = mock.Mock()
96        mock_ssh.GetBaseCmd.return_value = "/mock/ssh"
97        expected_image_shell_cmd = ("tar -cf - --lzop -S -C local/dir "
98                                    "super.img bootloader kernel android-info.txt | "
99                                    "/mock/ssh -- "
100                                    "tar -xf - --lzop -S -C remote/dir")
101        expected_target_files_shell_cmd = expected_image_shell_cmd.replace(
102            "local/dir", "local/dir/IMAGES")
103        expected_cvd_tar_ssh_cmd = "tar -xzf - -C remote/dir < /mock/cvd.tar.gz"
104        expected_cvd_dir_shell_cmd = ("tar -cf - --lzop -S -C /mock/cvd . | "
105                                      "/mock/ssh -- "
106                                      "tar -xf - --lzop -S -C remote/dir")
107
108        # Test with cvd directory.
109        mock_open = mock.mock_open(read_data="super.img\nbootloader\nkernel")
110        with mock.patch("acloud.internal.lib.cvd_utils.open", mock_open):
111            cvd_utils.UploadArtifacts(mock_ssh, "remote/dir","local/dir",
112                                      "/mock/cvd")
113        mock_open.assert_called_with("local/dir/required_images", "r",
114                                     encoding="utf-8")
115        mock_glob.glob.assert_called_once_with("local/dir/*.img")
116        mock_shell.assert_has_calls([mock.call(expected_image_shell_cmd),
117                                     mock.call(expected_cvd_dir_shell_cmd)])
118
119        # Test with required_images file.
120        mock_glob.glob.reset_mock()
121        mock_ssh.reset_mock()
122        mock_shell.reset_mock()
123        mock_open = mock.mock_open(read_data="super.img\nbootloader\nkernel")
124        with mock.patch("acloud.internal.lib.cvd_utils.open", mock_open):
125            cvd_utils.UploadArtifacts(mock_ssh, "remote/dir","local/dir",
126                                      "/mock/cvd.tar.gz")
127        mock_open.assert_called_with("local/dir/required_images", "r",
128                                     encoding="utf-8")
129        mock_glob.glob.assert_called_once_with("local/dir/*.img")
130        mock_shell.assert_called_with(expected_image_shell_cmd)
131        mock_ssh.Run.assert_called_with(expected_cvd_tar_ssh_cmd)
132
133        # Test with target files directory and glob.
134        mock_glob.glob.reset_mock()
135        mock_ssh.reset_mock()
136        mock_shell.reset_mock()
137        mock_glob.glob.side_effect = (
138            lambda path: [path.replace("*", "super")] if
139                         path.startswith("local/dir/IMAGES") else [])
140        with mock.patch("acloud.internal.lib.cvd_utils.open",
141                        side_effect=IOError("file does not exist")):
142            cvd_utils.UploadArtifacts(mock_ssh, "remote/dir", "local/dir",
143                                      "/mock/cvd.tar.gz")
144        self.assertGreater(mock_glob.glob.call_count, 2)
145        mock_shell.assert_called_with(expected_target_files_shell_cmd)
146        mock_ssh.Run.assert_called_with(expected_cvd_tar_ssh_cmd)
147
148    @mock.patch("acloud.internal.lib.cvd_utils.create_common")
149    def testUploadBootImages(self, mock_create_common):
150        """Test FindBootImages and UploadExtraImages."""
151        mock_ssh = mock.Mock()
152        with tempfile.TemporaryDirectory(prefix="cvd_utils") as image_dir:
153            mock_create_common.FindBootImage.return_value = "boot.img"
154            self.CreateFile(os.path.join(image_dir, "vendor_boot.img"))
155
156            mock_avd_spec = mock.Mock(local_kernel_image="boot.img",
157                                      local_system_image=None,
158                                      local_system_dlkm_image=None,
159                                      local_vendor_image=None,
160                                      local_vendor_boot_image=None)
161            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
162                                               None)
163            self.assertEqual([("-boot_image", "dir/acloud_image/boot.img")],
164                             args)
165            mock_ssh.Run.assert_called_once_with("mkdir -p dir/acloud_image")
166            mock_ssh.ScpPushFile.assert_called_once_with(
167                "boot.img", "dir/acloud_image/boot.img")
168
169            mock_ssh.reset_mock()
170            mock_avd_spec.local_kernel_image = image_dir
171            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
172                                               None)
173            self.assertEqual(
174                [("-boot_image", "dir/acloud_image/boot.img"),
175                 ("-vendor_boot_image", "dir/acloud_image/vendor_boot.img")],
176                args)
177            mock_ssh.Run.assert_called_once()
178            self.assertEqual(2, mock_ssh.ScpPushFile.call_count)
179
180    def testUploadKernelImages(self):
181        """Test FindKernelImages and UploadExtraImages."""
182        mock_ssh = mock.Mock()
183        with tempfile.TemporaryDirectory(prefix="cvd_utils") as image_dir:
184            kernel_image_path = os.path.join(image_dir, "Image")
185            self.CreateFile(kernel_image_path)
186            self.CreateFile(os.path.join(image_dir, "initramfs.img"))
187            self.CreateFile(os.path.join(image_dir, "boot.img"))
188
189            mock_avd_spec = mock.Mock(local_kernel_image=kernel_image_path,
190                                      local_system_image=None,
191                                      local_system_dlkm_image=None,
192                                      local_vendor_image=None,
193                                      local_vendor_boot_image=None)
194            with self.assertRaises(errors.GetLocalImageError):
195                cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
196                                            None)
197
198            mock_ssh.reset_mock()
199            mock_avd_spec.local_kernel_image = image_dir
200            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
201                                               None)
202            self.assertEqual(
203                [("-kernel_path", "dir/acloud_image/kernel"),
204                 ("-initramfs_path", "dir/acloud_image/initramfs.img")],
205                args)
206            mock_ssh.Run.assert_called_once()
207            self.assertEqual(2, mock_ssh.ScpPushFile.call_count)
208
209    @mock.patch("acloud.internal.lib.ota_tools.FindOtaTools")
210    @mock.patch("acloud.internal.lib.ssh.ShellCmdWithRetry")
211    def testUploadSuperImage(self, mock_shell, mock_find_ota_tools):
212        """Test UploadExtraImages."""
213        self.Patch(create_common, "GetNonEmptyEnvVars", return_value=[])
214        mock_ssh = mock.Mock()
215        mock_ota_tools_object = mock.Mock()
216        mock_find_ota_tools.return_value = mock_ota_tools_object
217
218        with tempfile.TemporaryDirectory(prefix="cvd_utils") as temp_dir:
219            target_files_dir = os.path.join(temp_dir, "target_files")
220            extra_image_dir = os.path.join(temp_dir, "extra")
221            mock_avd_spec = mock.Mock(local_kernel_image=None,
222                                      local_system_image=extra_image_dir,
223                                      local_system_dlkm_image=extra_image_dir,
224                                      local_vendor_image=extra_image_dir,
225                                      local_vendor_boot_image=None,
226                                      local_tool_dirs=[])
227            self.CreateFile(
228                os.path.join(target_files_dir, "IMAGES", "boot.img"))
229            self.CreateFile(
230                os.path.join(target_files_dir, "META", "misc_info.txt"))
231            for image_name in ["system.img", "system_dlkm.img", "vendor.img",
232                               "vendor_dlkm.img", "odm.img", "odm_dlkm.img"]:
233                self.CreateFile(os.path.join(extra_image_dir, image_name))
234            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
235                                               target_files_dir)
236
237        self.assertEqual(
238            [("-super_image", "dir/acloud_image/super.img"),
239             ("-vbmeta_image", "dir/acloud_image/vbmeta.img")],
240            args)
241        mock_find_ota_tools.assert_called_once_with([])
242        mock_ssh.Run.assert_called_once_with("mkdir -p dir/acloud_image")
243        # Super image
244        mock_shell.assert_called_once()
245        upload_args = mock_shell.call_args[0]
246        self.assertEqual(1, len(upload_args))
247        self.assertIn(" super.img", upload_args[0])
248        self.assertIn("dir/acloud_image", upload_args[0])
249        mock_ota_tools_object.MixSuperImage.assert_called_once_with(
250            mock.ANY, mock.ANY, os.path.join(target_files_dir, "IMAGES"),
251            system_image=os.path.join(extra_image_dir, "system.img"),
252            system_ext_image=None,
253            product_image=None,
254            system_dlkm_image=os.path.join(extra_image_dir, "system_dlkm.img"),
255            vendor_image=os.path.join(extra_image_dir, "vendor.img"),
256            vendor_dlkm_image=os.path.join(extra_image_dir, "vendor_dlkm.img"),
257            odm_image=os.path.join(extra_image_dir, "odm.img"),
258            odm_dlkm_image=os.path.join(extra_image_dir, "odm_dlkm.img"))
259        # vbmeta image
260        mock_ota_tools_object.MakeDisabledVbmetaImage.assert_called_once()
261        mock_ssh.ScpPushFile.assert_called_once_with(
262            mock.ANY, "dir/acloud_image/vbmeta.img")
263
264
265    def testUploadVendorBootImages(self):
266        """Test UploadExtraImages."""
267        mock_ssh = mock.Mock()
268        with tempfile.TemporaryDirectory(prefix="cvd_utils") as image_dir:
269            vendor_boot_image_path = os.path.join(image_dir,
270                                                  "vendor_boot-debug_test.img")
271            self.CreateFile(vendor_boot_image_path)
272
273            mock_avd_spec = mock.Mock(
274                local_kernel_image=None,
275                local_system_image=None,
276                local_system_dlkm_image=None,
277                local_vendor_image=None,
278                local_vendor_boot_image=vendor_boot_image_path)
279
280            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
281                                               None)
282            self.assertEqual(
283                [("-vendor_boot_image", "dir/acloud_image/vendor_boot.img")],
284                args)
285            mock_ssh.Run.assert_called_once()
286            mock_ssh.ScpPushFile.assert_called_once_with(
287                mock.ANY, "dir/acloud_image/vendor_boot.img")
288
289            mock_ssh.reset_mock()
290            self.CreateFile(os.path.join(image_dir, "vendor_boot.img"))
291            mock_avd_spec.local_vendor_boot_image = image_dir
292            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
293                                               None)
294            self.assertEqual(
295                [("-vendor_boot_image", "dir/acloud_image/vendor_boot.img")],
296                args)
297            mock_ssh.Run.assert_called_once()
298            mock_ssh.ScpPushFile.assert_called_once_with(
299                mock.ANY, "dir/acloud_image/vendor_boot.img")
300
301
302    def testCleanUpRemoteCvd(self):
303        """Test CleanUpRemoteCvd."""
304        mock_ssh = mock.Mock()
305        mock_ssh.Run.side_effect = ["", "", ""]
306        cvd_utils.CleanUpRemoteCvd(mock_ssh, "dir", raise_error=True)
307        mock_ssh.Run.assert_has_calls([
308            mock.call("'readlink -n -e dir/image_dir_link || true'"),
309            mock.call("'HOME=$HOME/dir dir/bin/stop_cvd'"),
310            mock.call("'rm -rf dir/*'")])
311
312        mock_ssh.reset_mock()
313        mock_ssh.Run.side_effect = ["img_dir", "", "", ""]
314        cvd_utils.CleanUpRemoteCvd(mock_ssh, "dir", raise_error=True)
315        mock_ssh.Run.assert_has_calls([
316            mock.call("'readlink -n -e dir/image_dir_link || true'"),
317            mock.call("'mkdir -p img_dir && flock img_dir.lock -c '\"'\"'"
318                      "rm -f dir/image_dir_link && "
319                      "expr $(test -s img_dir.lock && "
320                      "cat img_dir.lock || echo 1) - 1 > img_dir.lock || "
321                      "rm -rf img_dir img_dir.lock'\"'\"''"),
322            mock.call("'HOME=$HOME/dir dir/bin/stop_cvd'"),
323            mock.call("'rm -rf dir/*'")])
324
325        mock_ssh.reset_mock()
326        mock_ssh.Run.side_effect = [
327            "",
328            subprocess.CalledProcessError(cmd="should raise", returncode=1)]
329        with self.assertRaises(subprocess.CalledProcessError):
330            cvd_utils.CleanUpRemoteCvd(mock_ssh, "dir", raise_error=True)
331
332        mock_ssh.reset_mock()
333        mock_ssh.Run.side_effect = [
334            "",
335            subprocess.CalledProcessError(cmd="should ignore", returncode=1),
336            None]
337        cvd_utils.CleanUpRemoteCvd(mock_ssh, "dir", raise_error=False)
338        mock_ssh.Run.assert_any_call("'HOME=$HOME/dir dir/bin/stop_cvd'",
339                                     retry=0)
340        mock_ssh.Run.assert_any_call("'rm -rf dir/*'")
341
342    def testGetRemoteHostBaseDir(self):
343        """Test GetRemoteHostBaseDir."""
344        self.assertEqual("acloud_cf_1", cvd_utils.GetRemoteHostBaseDir(None))
345        self.assertEqual("acloud_cf_2", cvd_utils.GetRemoteHostBaseDir(2))
346
347    def testFormatRemoteHostInstanceName(self):
348        """Test FormatRemoteHostInstanceName."""
349        name = cvd_utils.FormatRemoteHostInstanceName(
350            self._REMOTE_HOST_IP, None, self._BUILD_ID, self._PRODUCT_NAME)
351        self.assertEqual(name, self._REMOTE_HOST_INSTANCE_NAME_1)
352
353        name = cvd_utils.FormatRemoteHostInstanceName(
354            self._REMOTE_HOST_IP, 2, self._BUILD_ID, self._PRODUCT_NAME)
355        self.assertEqual(name, self._REMOTE_HOST_INSTANCE_NAME_2)
356
357    def testParseRemoteHostAddress(self):
358        """Test ParseRemoteHostAddress."""
359        result = cvd_utils.ParseRemoteHostAddress(
360            self._REMOTE_HOST_INSTANCE_NAME_1)
361        self.assertEqual(result, (self._REMOTE_HOST_IP, "acloud_cf_1"))
362
363        result = cvd_utils.ParseRemoteHostAddress(
364            self._REMOTE_HOST_INSTANCE_NAME_2)
365        self.assertEqual(result, (self._REMOTE_HOST_IP, "acloud_cf_2"))
366
367        result = cvd_utils.ParseRemoteHostAddress(
368            "host-goldfish-192.0.2.1-5554-123456-sdk_x86_64-sdk")
369        self.assertIsNone(result)
370
371    # pylint: disable=protected-access
372    def testRemoteImageDirLink(self):
373        """Test PrepareRemoteImageDirLink and _DeleteRemoteImageDirLink."""
374        self.assertEqual(os.path, cvd_utils.remote_path)
375        with tempfile.TemporaryDirectory(prefix="cvd_utils") as temp_dir:
376            env = os.environ.copy()
377            env["HOME"] = temp_dir
378            # Execute the commands locally.
379            mock_ssh = mock.Mock()
380            mock_ssh.Run.side_effect = lambda cmd: subprocess.check_output(
381                "sh -c " + cmd, shell=True, cwd=temp_dir, env=env
382            ).decode("utf-8")
383            # Relative paths under temp_dir.
384            base_dir_name_1 = "acloud_cf_1"
385            base_dir_name_2 = "acloud_cf_2"
386            image_dir_name = "test/img"
387            rel_ref_cnt_path = "test/img.lock"
388            # Absolute paths.
389            image_dir = os.path.join(temp_dir, image_dir_name)
390            ref_cnt_path = os.path.join(temp_dir, rel_ref_cnt_path)
391            link_path_1 = os.path.join(temp_dir, base_dir_name_1,
392                                       "image_dir_link")
393            link_path_2 = os.path.join(temp_dir, base_dir_name_2,
394                                       "image_dir_link")
395            # Delete non-existing directories.
396            cvd_utils._DeleteRemoteImageDirLink(mock_ssh, base_dir_name_1)
397            mock_ssh.Run.assert_called_with(
398                f"'readlink -n -e {base_dir_name_1}/image_dir_link || true'")
399            self.assertFalse(
400                os.path.exists(os.path.join(temp_dir, base_dir_name_1)))
401            self.assertFalse(os.path.exists(image_dir))
402            self.assertFalse(os.path.exists(ref_cnt_path))
403            # Prepare the first base dir.
404            cvd_utils.PrepareRemoteImageDirLink(mock_ssh, base_dir_name_1,
405                                                image_dir_name)
406            mock_ssh.Run.assert_called_with(
407                f"'mkdir -p {image_dir_name} && flock {rel_ref_cnt_path} -c "
408                f"'\"'\"'mkdir -p {base_dir_name_1} {image_dir_name} && "
409                f"ln -s -r {image_dir_name} "
410                f"{base_dir_name_1}/image_dir_link && "
411                f"expr $(test -s {rel_ref_cnt_path} && "
412                f"cat {rel_ref_cnt_path} || echo 0) + 1 > "
413                f"{rel_ref_cnt_path}'\"'\"''")
414            self.assertTrue(os.path.islink(link_path_1))
415            self.assertEqual("../test/img", os.readlink(link_path_1))
416            self.assertTrue(os.path.isfile(ref_cnt_path))
417            with open(ref_cnt_path, "r", encoding="utf-8") as ref_cnt_file:
418                self.assertEqual("1\n", ref_cnt_file.read())
419            # Prepare the second base dir.
420            cvd_utils.PrepareRemoteImageDirLink(mock_ssh, base_dir_name_2,
421                                                image_dir_name)
422            self.assertTrue(os.path.islink(link_path_2))
423            self.assertEqual("../test/img", os.readlink(link_path_2))
424            self.assertTrue(os.path.isfile(ref_cnt_path))
425            with open(ref_cnt_path, "r", encoding="utf-8") as ref_cnt_file:
426                self.assertEqual("2\n", ref_cnt_file.read())
427            # Delete the first base dir.
428            cvd_utils._DeleteRemoteImageDirLink(mock_ssh, base_dir_name_1)
429            self.assertFalse(os.path.lexists(link_path_1))
430            self.assertTrue(os.path.isfile(ref_cnt_path))
431            with open(ref_cnt_path, "r", encoding="utf-8") as ref_cnt_file:
432                self.assertEqual("1\n", ref_cnt_file.read())
433            # Delete the second base dir.
434            cvd_utils._DeleteRemoteImageDirLink(mock_ssh, base_dir_name_2)
435            self.assertFalse(os.path.lexists(link_path_2))
436            self.assertFalse(os.path.exists(image_dir))
437            self.assertFalse(os.path.exists(ref_cnt_path))
438
439    @mock.patch("acloud.internal.lib.cvd_utils.utils.PollAndWait")
440    @mock.patch("acloud.internal.lib.cvd_utils.utils.time.time",
441                return_value=90.0)
442    def testLoadRemoteImageArgs(self, _mock_time, mock_poll_and_wait):
443        """Test LoadRemoteImageArgs."""
444        deadline = 99.9
445        self.assertEqual(os.path, cvd_utils.remote_path)
446
447        with tempfile.TemporaryDirectory(prefix="cvd_utils") as temp_dir:
448            env = os.environ.copy()
449            env["HOME"] = temp_dir
450            # Execute the commands locally.
451            mock_ssh = mock.Mock()
452            mock_ssh.Run.side_effect = lambda cmd: subprocess.check_output(
453                "sh -c " + cmd, shell=True, cwd=temp_dir, env=env, text=True)
454            mock_poll_and_wait.side_effect = lambda func, **kwargs: func()
455
456            timestamp_path = os.path.join(temp_dir, "timestamp.txt")
457            args_path = os.path.join(temp_dir, "args.txt")
458
459            # Test with an uninitialized directory.
460            args = cvd_utils.LoadRemoteImageArgs(
461                mock_ssh, timestamp_path, args_path, deadline)
462
463            self.assertIsNone(args)
464            mock_ssh.Run.assert_called_once()
465            with open(timestamp_path, "r", encoding="utf-8") as timestamp_file:
466                timestamp = timestamp_file.read().strip()
467                self.assertRegex(timestamp, r"\d+",
468                                 f"Invalid timestamp: {timestamp}")
469            self.assertFalse(os.path.exists(args_path))
470
471            # Test with an initialized directory and the uploader times out.
472            mock_ssh.Run.reset_mock()
473
474            with self.assertRaises(errors.CreateError):
475                cvd_utils.LoadRemoteImageArgs(
476                    mock_ssh, timestamp_path, args_path, deadline)
477
478            mock_ssh.Run.assert_has_calls([
479                mock.call(f"'flock {timestamp_path} -c '\"'\"'"
480                          f"test -s {timestamp_path} && "
481                          f"cat {timestamp_path} || "
482                          f"expr $(date +%s) + 9 > {timestamp_path}'\"'\"''"),
483                mock.call(f"'flock {args_path} -c '\"'\"'"
484                          f"test -s {args_path} -o "
485                          f"{timestamp} -le $(date +%s) || "
486                          "echo wait...'\"'\"''"),
487                mock.call(f"'flock {args_path} -c '\"'\"'"
488                          f"cat {args_path}'\"'\"''")
489            ])
490            with open(timestamp_path, "r", encoding="utf-8") as timestamp_file:
491                self.assertEqual(timestamp_file.read().strip(), timestamp)
492            self.assertEqual(os.path.getsize(args_path), 0)
493
494            # Test with an initialized directory.
495            mock_ssh.Run.reset_mock()
496            self.CreateFile(args_path, b'[["arg", "1"]]')
497
498            args = cvd_utils.LoadRemoteImageArgs(
499                mock_ssh, timestamp_path, args_path, deadline)
500
501            self.assertEqual(args, [["arg", "1"]])
502            self.assertEqual(mock_ssh.Run.call_count, 3)
503
504    def testSaveRemoteImageArgs(self):
505        """Test SaveRemoteImageArgs."""
506        with tempfile.TemporaryDirectory(prefix="cvd_utils") as temp_dir:
507            env = os.environ.copy()
508            env["HOME"] = temp_dir
509            mock_ssh = mock.Mock()
510            mock_ssh.Run.side_effect = lambda cmd: subprocess.check_call(
511                "sh -c " + cmd, shell=True, cwd=temp_dir, env=env, text=True)
512            args_path = os.path.join(temp_dir, "args.txt")
513
514            cvd_utils.SaveRemoteImageArgs(mock_ssh, args_path, [("arg", "1")])
515
516            mock_ssh.Run.assert_called_with(
517                f"'flock {args_path} -c '\"'\"'"
518                f"""echo '"'"'"'"'"'"'"'"'[["arg", "1"]]'"'"'"'"'"'"'"'"' > """
519                f"{args_path}'\"'\"''")
520            with open(args_path, "r", encoding="utf-8") as args_file:
521                self.assertEqual(args_file.read().strip(), '[["arg", "1"]]')
522
523    def testGetConfigFromRemoteAndroidInfo(self):
524        """Test GetConfigFromRemoteAndroidInfo."""
525        mock_ssh = mock.Mock()
526        mock_ssh.GetCmdOutput.return_value = "require board=vsoc_x86_64\n"
527        config = cvd_utils.GetConfigFromRemoteAndroidInfo(mock_ssh, ".")
528        mock_ssh.GetCmdOutput.assert_called_with("cat ./android-info.txt")
529        self.assertIsNone(config)
530
531        mock_ssh.GetCmdOutput.return_value += "config=phone\n"
532        config = cvd_utils.GetConfigFromRemoteAndroidInfo(mock_ssh, ".")
533        self.assertEqual(config, "phone")
534
535    def testGetRemoteLaunchCvdCmd(self):
536        """Test GetRemoteLaunchCvdCmd."""
537        # Minimum arguments
538        mock_cfg = mock.Mock(extra_data_disk_size_gb=0)
539        hw_property = {
540            constants.HW_X_RES: "1080",
541            constants.HW_Y_RES: "1920",
542            constants.HW_ALIAS_DPI: "240"}
543        mock_avd_spec = mock.Mock(
544            spec=[],
545            cfg=mock_cfg,
546            hw_customize=False,
547            hw_property=hw_property,
548            connect_webrtc=False,
549            connect_vnc=False,
550            openwrt=False,
551            num_avds_per_instance=1,
552            base_instance_num=0,
553            launch_args="")
554        expected_cmd = (
555            "HOME=$HOME/dir dir/bin/launch_cvd -daemon "
556            "-x_res=1080 -y_res=1920 -dpi=240 "
557            "-undefok=report_anonymous_usage_stats,config "
558            "-report_anonymous_usage_stats=y")
559        cmd = cvd_utils.GetRemoteLaunchCvdCmd("dir", mock_avd_spec,
560                                              config=None, extra_args=())
561        self.assertEqual(cmd, expected_cmd)
562
563        # All arguments.
564        mock_cfg = mock.Mock(extra_data_disk_size_gb=20)
565        hw_property = {
566            constants.HW_X_RES: "1080",
567            constants.HW_Y_RES: "1920",
568            constants.HW_ALIAS_DPI: "240",
569            constants.HW_ALIAS_DISK: "10240",
570            constants.HW_ALIAS_CPUS: "2",
571            constants.HW_ALIAS_MEMORY: "4096"}
572        mock_avd_spec = mock.Mock(
573            spec=[],
574            cfg=mock_cfg,
575            hw_customize=True,
576            hw_property=hw_property,
577            connect_webrtc=True,
578            webrtc_device_id="pet-name",
579            connect_vnc=True,
580            openwrt=True,
581            num_avds_per_instance=2,
582            base_instance_num=3,
583            launch_args="--setupwizard_mode=REQUIRED")
584        expected_cmd = (
585            "HOME=$HOME/dir dir/bin/launch_cvd -daemon --extra args "
586            "-data_policy=create_if_missing -blank_data_image_mb=20480 "
587            "-config=phone -x_res=1080 -y_res=1920 -dpi=240 "
588            "-data_policy=always_create -blank_data_image_mb=10240 "
589            "-cpus=2 -memory_mb=4096 "
590            "--start_webrtc --vm_manager=crosvm "
591            "--webrtc_device_id=pet-name "
592            "--start_vnc_server=true "
593            "-console=true "
594            "-num_instances=2 --base_instance_num=3 "
595            "--setupwizard_mode=REQUIRED "
596            "-undefok=report_anonymous_usage_stats,config "
597            "-report_anonymous_usage_stats=y")
598        cmd = cvd_utils.GetRemoteLaunchCvdCmd(
599            "dir", mock_avd_spec, "phone", ("--extra", "args"))
600        self.assertEqual(cmd, expected_cmd)
601
602    def testExecuteRemoteLaunchCvd(self):
603        """Test ExecuteRemoteLaunchCvd."""
604        mock_ssh = mock.Mock()
605        error_msg = cvd_utils.ExecuteRemoteLaunchCvd(mock_ssh, "launch_cvd", 1)
606        self.assertFalse(error_msg)
607        mock_ssh.Run.assert_called()
608
609        mock_ssh.Run.side_effect = errors.LaunchCVDFail(
610            "Test unknown command line flag 'start_vnc_server'.")
611        error_msg = cvd_utils.ExecuteRemoteLaunchCvd(mock_ssh, "launch_cvd", 1)
612        self.assertIn("VNC is not supported in the current build.", error_msg)
613
614    def testGetRemoteFetcherConfigJson(self):
615        """Test GetRemoteFetcherConfigJson."""
616        expected_log = {"path": "dir/fetcher_config.json",
617                        "type": constants.LOG_TYPE_CUTTLEFISH_LOG}
618        self.assertEqual(expected_log,
619                         cvd_utils.GetRemoteFetcherConfigJson("dir"))
620
621    @mock.patch("acloud.internal.lib.cvd_utils.utils")
622    def testFindRemoteLogs(self, mock_utils):
623        """Test FindRemoteLogs with the runtime directories in Android 13."""
624        mock_ssh = mock.Mock()
625        mock_utils.FindRemoteFiles.return_value = [
626            "/kernel.log", "/logcat", "/launcher.log", "/access-kregistry",
627            "/cuttlefish_config.json"]
628
629        logs = cvd_utils.FindRemoteLogs(mock_ssh, "dir", None, None)
630        mock_ssh.Run.assert_called_with(
631            "test -d dir/cuttlefish/instances/cvd-1", retry=0)
632        mock_utils.FindRemoteFiles.assert_called_with(
633            mock_ssh, ["dir/cuttlefish/instances/cvd-1"])
634        expected_logs = [
635            {
636                "path": "/kernel.log",
637                "type": constants.LOG_TYPE_KERNEL_LOG,
638                "name": "kernel.log"
639            },
640            {
641                "path": "/logcat",
642                "type": constants.LOG_TYPE_LOGCAT,
643                "name": "full_gce_logcat"
644            },
645            {
646                "path": "/launcher.log",
647                "type": constants.LOG_TYPE_CUTTLEFISH_LOG,
648                "name": "launcher.log"
649            },
650            {
651                "path": "/cuttlefish_config.json",
652                "type": constants.LOG_TYPE_CUTTLEFISH_LOG,
653                "name": "cuttlefish_config.json"
654            },
655            {
656                "path": "dir/cuttlefish/instances/cvd-1/tombstones",
657                "type": constants.LOG_TYPE_DIR,
658                "name": "tombstones-zip"
659            },
660        ]
661        self.assertEqual(expected_logs, logs)
662
663    @mock.patch("acloud.internal.lib.cvd_utils.utils")
664    def testFindRemoteLogsWithLegacyDirs(self, mock_utils):
665        """Test FindRemoteLogs with the runtime directories in Android 11."""
666        mock_ssh = mock.Mock()
667        mock_ssh.Run.side_effect = subprocess.CalledProcessError(
668            cmd="test", returncode=1)
669        mock_utils.FindRemoteFiles.return_value = [
670            "dir/cuttlefish_runtime/kernel.log",
671            "dir/cuttlefish_runtime.4/kernel.log",
672        ]
673
674        logs = cvd_utils.FindRemoteLogs(mock_ssh, "dir", 3, 2)
675        mock_ssh.Run.assert_called_with(
676            "test -d dir/cuttlefish/instances/cvd-3", retry=0)
677        mock_utils.FindRemoteFiles.assert_called_with(
678            mock_ssh, ["dir/cuttlefish_runtime", "dir/cuttlefish_runtime.4"])
679        expected_logs = [
680            {
681                "path": "dir/cuttlefish_runtime/kernel.log",
682                "type": constants.LOG_TYPE_KERNEL_LOG,
683                "name": "kernel.log"
684            },
685            {
686                "path": "dir/cuttlefish_runtime.4/kernel.log",
687                "type": constants.LOG_TYPE_KERNEL_LOG,
688                "name": "kernel.1.log"
689            },
690            {
691                "path": "dir/cuttlefish_runtime/tombstones",
692                "type": constants.LOG_TYPE_DIR,
693                "name": "tombstones-zip"
694            },
695            {
696                "path": "dir/cuttlefish_runtime.4/tombstones",
697                "type": constants.LOG_TYPE_DIR,
698                "name": "tombstones-zip.1"
699            },
700        ]
701        self.assertEqual(expected_logs, logs)
702
703    def testFindLocalLogs(self):
704        """Test FindLocalLogs with the runtime directory in Android 13."""
705        with tempfile.TemporaryDirectory() as temp_dir:
706            log_dir = os.path.join(temp_dir, "instances", "cvd-2", "logs")
707            kernel_log = os.path.join(os.path.join(log_dir, "kernel.log"))
708            launcher_log = os.path.join(os.path.join(log_dir, "launcher.log"))
709            logcat = os.path.join(os.path.join(log_dir, "logcat"))
710            self.CreateFile(kernel_log)
711            self.CreateFile(launcher_log)
712            self.CreateFile(logcat)
713            self.CreateFile(os.path.join(temp_dir, "legacy.log"))
714            self.CreateFile(os.path.join(log_dir, "log.txt"))
715            os.symlink(os.path.join(log_dir, "launcher.log"),
716                       os.path.join(log_dir, "link.log"))
717
718            logs = cvd_utils.FindLocalLogs(temp_dir, 2)
719            expected_logs = [
720                {
721                    "path": kernel_log,
722                    "type": constants.LOG_TYPE_KERNEL_LOG,
723                },
724                {
725                    "path": launcher_log,
726                    "type": constants.LOG_TYPE_CUTTLEFISH_LOG,
727                },
728                {
729                    "path": logcat,
730                    "type": constants.LOG_TYPE_LOGCAT,
731                },
732            ]
733            self.assertEqual(expected_logs,
734                             sorted(logs, key=lambda log: log["path"]))
735
736    def testFindLocalLogsWithLegacyDir(self):
737        """Test FindLocalLogs with the runtime directory in Android 11."""
738        with tempfile.TemporaryDirectory() as temp_dir:
739            log_dir = os.path.join(temp_dir, "cuttlefish_runtime.2")
740            log_dir_link = os.path.join(temp_dir, "cuttlefish_runtime")
741            os.mkdir(log_dir)
742            os.symlink(log_dir, log_dir_link, target_is_directory=True)
743            launcher_log = os.path.join(log_dir_link, "launcher.log")
744            self.CreateFile(launcher_log)
745
746            logs = cvd_utils.FindLocalLogs(log_dir_link, 2)
747            expected_logs = [
748                {
749                    "path": launcher_log,
750                    "type": constants.LOG_TYPE_CUTTLEFISH_LOG,
751                },
752            ]
753            self.assertEqual(expected_logs, logs)
754
755    def testGetOpenWrtInfoDict(self):
756        """Test GetOpenWrtInfoDict."""
757        mock_ssh = mock.Mock()
758        mock_ssh.GetBaseCmd.return_value = "/mock/ssh"
759        openwrt_info = {
760            "ssh_command": "/mock/ssh",
761            "screen_command": "screen ./cuttlefish_runtime/console"}
762        self.assertDictEqual(openwrt_info,
763                             cvd_utils.GetOpenWrtInfoDict(mock_ssh, "."))
764        mock_ssh.GetBaseCmd.assert_called_with("ssh")
765
766    def testGetRemoteBuildInfoDict(self):
767        """Test GetRemoteBuildInfoDict."""
768        remote_image = {
769            "branch": "aosp-android-12-gsi",
770            "build_id": "100000",
771            "build_target": "aosp_cf_x86_64_phone-userdebug"}
772        mock_avd_spec = mock.Mock(
773            spec=[],
774            remote_image=remote_image,
775            kernel_build_info={"build_target": "kernel"},
776            system_build_info={},
777            bootloader_build_info={},
778            android_efi_loader_build_info = {})
779        self.assertEqual(remote_image,
780                         cvd_utils.GetRemoteBuildInfoDict(mock_avd_spec))
781
782        kernel_build_info = {
783            "branch": "aosp_kernel-common-android12-5.10",
784            "build_id": "200000",
785            "build_target": "kernel_virt_x86_64"}
786        system_build_info = {
787            "branch": "aosp-android-12-gsi",
788            "build_id": "300000",
789            "build_target": "aosp_x86_64-userdebug"}
790        bootloader_build_info = {
791            "branch": "aosp_u-boot-mainline",
792            "build_id": "400000",
793            "build_target": "u-boot_crosvm_x86_64"}
794        android_efi_loader_build_info = {
795            "build_id": "500000",
796            "artifact": "gbl_aarch64.efi"
797        }
798        all_build_info = {
799            "kernel_branch": "aosp_kernel-common-android12-5.10",
800            "kernel_build_id": "200000",
801            "kernel_build_target": "kernel_virt_x86_64",
802            "system_branch": "aosp-android-12-gsi",
803            "system_build_id": "300000",
804            "system_build_target": "aosp_x86_64-userdebug",
805            "bootloader_branch": "aosp_u-boot-mainline",
806            "bootloader_build_id": "400000",
807            "bootloader_build_target": "u-boot_crosvm_x86_64",
808            "android_efi_loader_build_id": "500000",
809            "android_efi_loader_artifact": "gbl_aarch64.efi"
810        }
811        all_build_info.update(remote_image)
812        mock_avd_spec = mock.Mock(
813            spec=[],
814            remote_image=remote_image,
815            kernel_build_info=kernel_build_info,
816            system_build_info=system_build_info,
817            bootloader_build_info=bootloader_build_info,
818            android_efi_loader_build_info=android_efi_loader_build_info)
819        self.assertEqual(all_build_info,
820                         cvd_utils.GetRemoteBuildInfoDict(mock_avd_spec))
821
822    def testFindMiscInfo(self):
823        """Test FindMiscInfo."""
824        with tempfile.TemporaryDirectory() as temp_dir:
825            with self.assertRaises(errors.CheckPathError):
826                cvd_utils.FindMiscInfo(temp_dir)
827            misc_info_path = os.path.join(temp_dir, "META", "misc_info.txt")
828            self.CreateFile(misc_info_path, b"key=value")
829            self.assertEqual(misc_info_path, cvd_utils.FindMiscInfo(temp_dir))
830
831    def testFindImageDir(self):
832        """Test FindImageDir."""
833        with tempfile.TemporaryDirectory() as temp_dir:
834            with self.assertRaises(errors.GetLocalImageError):
835                cvd_utils.FindImageDir(temp_dir)
836            image_dir = os.path.join(temp_dir, "IMAGES")
837            self.CreateFile(os.path.join(image_dir, "super.img"))
838            self.assertEqual(image_dir, cvd_utils.FindImageDir(temp_dir))
839
840
841if __name__ == "__main__":
842    unittest.main()
843