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