xref: /aosp_15_r20/tools/acloud/pull/pull_test.py (revision 800a58d989c669b8eb8a71d8df53b1ba3d411444)
1*800a58d9SAndroid Build Coastguard Worker# Copyright 2019 - The Android Open Source Project
2*800a58d9SAndroid Build Coastguard Worker#
3*800a58d9SAndroid Build Coastguard Worker# Licensed under the Apache License, Version 2.0 (the "License");
4*800a58d9SAndroid Build Coastguard Worker# you may not use this file except in compliance with the License.
5*800a58d9SAndroid Build Coastguard Worker# You may obtain a copy of the License at
6*800a58d9SAndroid Build Coastguard Worker#
7*800a58d9SAndroid Build Coastguard Worker#     http://www.apache.org/licenses/LICENSE-2.0
8*800a58d9SAndroid Build Coastguard Worker#
9*800a58d9SAndroid Build Coastguard Worker# Unless required by applicable law or agreed to in writing, software
10*800a58d9SAndroid Build Coastguard Worker# distributed under the License is distributed on an "AS IS" BASIS,
11*800a58d9SAndroid Build Coastguard Worker# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12*800a58d9SAndroid Build Coastguard Worker# See the License for the specific language governing permissions and
13*800a58d9SAndroid Build Coastguard Worker# limitations under the License.
14*800a58d9SAndroid Build Coastguard Worker"""Tests for pull."""
15*800a58d9SAndroid Build Coastguard Workerimport unittest
16*800a58d9SAndroid Build Coastguard Worker
17*800a58d9SAndroid Build Coastguard Workerimport os
18*800a58d9SAndroid Build Coastguard Workerimport tempfile
19*800a58d9SAndroid Build Coastguard Worker
20*800a58d9SAndroid Build Coastguard Workerfrom unittest import mock
21*800a58d9SAndroid Build Coastguard Worker
22*800a58d9SAndroid Build Coastguard Workerfrom acloud import errors
23*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal import constants
24*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal.lib import driver_test_lib
25*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal.lib import ssh
26*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal.lib import utils
27*800a58d9SAndroid Build Coastguard Workerfrom acloud.list import list as list_instances
28*800a58d9SAndroid Build Coastguard Workerfrom acloud.public import config
29*800a58d9SAndroid Build Coastguard Workerfrom acloud.pull import pull
30*800a58d9SAndroid Build Coastguard Worker
31*800a58d9SAndroid Build Coastguard Worker
32*800a58d9SAndroid Build Coastguard Workerclass PullTest(driver_test_lib.BaseDriverTest):
33*800a58d9SAndroid Build Coastguard Worker    """Test pull."""
34*800a58d9SAndroid Build Coastguard Worker
35*800a58d9SAndroid Build Coastguard Worker    # pylint: disable=no-member
36*800a58d9SAndroid Build Coastguard Worker    def testPullFileFromInstance(self):
37*800a58d9SAndroid Build Coastguard Worker        """test PullFileFromInstance."""
38*800a58d9SAndroid Build Coastguard Worker        cfg = mock.MagicMock()
39*800a58d9SAndroid Build Coastguard Worker        cfg.ssh_private_key_path = "fake_ssh_path"
40*800a58d9SAndroid Build Coastguard Worker        cfg.extra_args_ssh_tunnel = ""
41*800a58d9SAndroid Build Coastguard Worker        instance = mock.MagicMock()
42*800a58d9SAndroid Build Coastguard Worker        instance.ip = "1.1.1.1"
43*800a58d9SAndroid Build Coastguard Worker        # Multiple selected files case.
44*800a58d9SAndroid Build Coastguard Worker        selected_files = ["file1.log", "file2.log"]
45*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "SelectLogFileToPull", return_value=selected_files)
46*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "PullLogs")
47*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "DisplayLog")
48*800a58d9SAndroid Build Coastguard Worker        pull.PullFileFromInstance(cfg, instance)
49*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(pull.DisplayLog.call_count, 0)
50*800a58d9SAndroid Build Coastguard Worker
51*800a58d9SAndroid Build Coastguard Worker        # Only one file selected case.
52*800a58d9SAndroid Build Coastguard Worker        selected_files = ["file1.log"]
53*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "SelectLogFileToPull", return_value=selected_files)
54*800a58d9SAndroid Build Coastguard Worker        pull.PullFileFromInstance(cfg, instance)
55*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(pull.DisplayLog.call_count, 1)
56*800a58d9SAndroid Build Coastguard Worker
57*800a58d9SAndroid Build Coastguard Worker    def testPullLogs(self):
58*800a58d9SAndroid Build Coastguard Worker        """test PullLogs."""
59*800a58d9SAndroid Build Coastguard Worker        self.Patch(tempfile, "gettempdir", return_value="/tmp")
60*800a58d9SAndroid Build Coastguard Worker        self.Patch(os.path, "exists", return_value=False)
61*800a58d9SAndroid Build Coastguard Worker        mock_makedirs = self.Patch(os, "makedirs")
62*800a58d9SAndroid Build Coastguard Worker        _ssh = mock.MagicMock()
63*800a58d9SAndroid Build Coastguard Worker        self.Patch(utils, "PrintColorString")
64*800a58d9SAndroid Build Coastguard Worker
65*800a58d9SAndroid Build Coastguard Worker        log_files = ["file1.log", "file2.log"]
66*800a58d9SAndroid Build Coastguard Worker        download_folder = pull.PullLogs(_ssh, log_files, "instance")
67*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(download_folder, "/tmp/instance")
68*800a58d9SAndroid Build Coastguard Worker        mock_makedirs.assert_called_once_with("/tmp/instance")
69*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(_ssh.ScpPullFile.call_count, 2)
70*800a58d9SAndroid Build Coastguard Worker        utils.PrintColorString.assert_called_once()
71*800a58d9SAndroid Build Coastguard Worker
72*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(ssh.Ssh, "Run")
73*800a58d9SAndroid Build Coastguard Worker    def testDisplayLog(self, mock_ssh_run):
74*800a58d9SAndroid Build Coastguard Worker        """Test DisplayLog."""
75*800a58d9SAndroid Build Coastguard Worker        fake_ip = ssh.IP(external="1.1.1.1", internal="10.1.1.1")
76*800a58d9SAndroid Build Coastguard Worker        _ssh = ssh.Ssh(ip=fake_ip,
77*800a58d9SAndroid Build Coastguard Worker                       user=constants.GCE_USER,
78*800a58d9SAndroid Build Coastguard Worker                       ssh_private_key_path="/fake/acloud_rea")
79*800a58d9SAndroid Build Coastguard Worker        self.Patch(utils, "GetUserAnswerYes", return_value="Y")
80*800a58d9SAndroid Build Coastguard Worker        log_file = "file1.log"
81*800a58d9SAndroid Build Coastguard Worker        pull.DisplayLog(_ssh, log_file)
82*800a58d9SAndroid Build Coastguard Worker        expected_cmd = "tail -f -n +1 %s" % log_file
83*800a58d9SAndroid Build Coastguard Worker        mock_ssh_run.assert_has_calls([
84*800a58d9SAndroid Build Coastguard Worker            mock.call(expected_cmd, show_output=True)])
85*800a58d9SAndroid Build Coastguard Worker
86*800a58d9SAndroid Build Coastguard Worker    def testSelectLogFileToPull(self):
87*800a58d9SAndroid Build Coastguard Worker        """test choose log files from the remote instance."""
88*800a58d9SAndroid Build Coastguard Worker        _ssh = mock.MagicMock()
89*800a58d9SAndroid Build Coastguard Worker
90*800a58d9SAndroid Build Coastguard Worker        # Test only one log file case
91*800a58d9SAndroid Build Coastguard Worker        log_files = ["file1.log"]
92*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "GetAllLogFilePaths", return_value=log_files)
93*800a58d9SAndroid Build Coastguard Worker        expected_result = ["file1.log"]
94*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(pull.SelectLogFileToPull(_ssh), expected_result)
95*800a58d9SAndroid Build Coastguard Worker
96*800a58d9SAndroid Build Coastguard Worker        # Test no log files case
97*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "GetAllLogFilePaths", return_value=[])
98*800a58d9SAndroid Build Coastguard Worker        with self.assertRaises(errors.CheckPathError):
99*800a58d9SAndroid Build Coastguard Worker            pull.SelectLogFileToPull(_ssh)
100*800a58d9SAndroid Build Coastguard Worker
101*800a58d9SAndroid Build Coastguard Worker        # Test two log files case.
102*800a58d9SAndroid Build Coastguard Worker        log_files = ["file1.log", "file2.log"]
103*800a58d9SAndroid Build Coastguard Worker        choose_log = ["file2.log"]
104*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "GetAllLogFilePaths", return_value=log_files)
105*800a58d9SAndroid Build Coastguard Worker        self.Patch(utils, "GetAnswerFromList", return_value=choose_log)
106*800a58d9SAndroid Build Coastguard Worker        expected_result = ["file2.log"]
107*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(pull.SelectLogFileToPull(_ssh), expected_result)
108*800a58d9SAndroid Build Coastguard Worker
109*800a58d9SAndroid Build Coastguard Worker        # Test user provided file name exist.
110*800a58d9SAndroid Build Coastguard Worker        log_files = ["cuttlefish_runtime/file1.log",
111*800a58d9SAndroid Build Coastguard Worker                     "cuttlefish_runtime/file2.log"]
112*800a58d9SAndroid Build Coastguard Worker        input_file = "file1.log"
113*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "GetAllLogFilePaths", return_value=log_files)
114*800a58d9SAndroid Build Coastguard Worker        expected_result = ["cuttlefish_runtime/file1.log"]
115*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(pull.SelectLogFileToPull(_ssh, input_file), expected_result)
116*800a58d9SAndroid Build Coastguard Worker
117*800a58d9SAndroid Build Coastguard Worker        # Test user provided file name not exist.
118*800a58d9SAndroid Build Coastguard Worker        log_files = ["cuttlefish_runtime/file1.log",
119*800a58d9SAndroid Build Coastguard Worker                     "cuttlefish_runtime/file2.log"]
120*800a58d9SAndroid Build Coastguard Worker        input_file = "not_exist.log"
121*800a58d9SAndroid Build Coastguard Worker        self.Patch(pull, "GetAllLogFilePaths", return_value=log_files)
122*800a58d9SAndroid Build Coastguard Worker        with self.assertRaises(errors.CheckPathError):
123*800a58d9SAndroid Build Coastguard Worker            pull.SelectLogFileToPull(_ssh, input_file)
124*800a58d9SAndroid Build Coastguard Worker
125*800a58d9SAndroid Build Coastguard Worker    def testGetAllLogFilePaths(self):
126*800a58d9SAndroid Build Coastguard Worker        """test that GetAllLogFilePaths can filter logs."""
127*800a58d9SAndroid Build Coastguard Worker        mock_find = self.Patch(utils, "FindRemoteFiles",
128*800a58d9SAndroid Build Coastguard Worker                               return_value=["kernel.log", "logcat", "kernel"])
129*800a58d9SAndroid Build Coastguard Worker        # Filter out file name is "kernel".
130*800a58d9SAndroid Build Coastguard Worker        expected_result = ["kernel.log", "logcat"]
131*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(pull.GetAllLogFilePaths(mock.Mock(), "unit/test"),
132*800a58d9SAndroid Build Coastguard Worker                         expected_result)
133*800a58d9SAndroid Build Coastguard Worker        mock_find.assert_called_with(mock.ANY, ["unit/test"])
134*800a58d9SAndroid Build Coastguard Worker
135*800a58d9SAndroid Build Coastguard Worker        # Filter out file extension is ".img".
136*800a58d9SAndroid Build Coastguard Worker        mock_find.return_value = ["kernel.log", "system.img", "userdata.img",
137*800a58d9SAndroid Build Coastguard Worker                                  "launcher.log"]
138*800a58d9SAndroid Build Coastguard Worker        expected_result = ["kernel.log", "launcher.log"]
139*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(pull.GetAllLogFilePaths(mock.Mock(), "unit/test"),
140*800a58d9SAndroid Build Coastguard Worker                         expected_result)
141*800a58d9SAndroid Build Coastguard Worker
142*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(pull, "PullFileFromInstance")
143*800a58d9SAndroid Build Coastguard Worker    def testRun(self, mock_pull_file):
144*800a58d9SAndroid Build Coastguard Worker        """test Run."""
145*800a58d9SAndroid Build Coastguard Worker        cfg = mock.MagicMock()
146*800a58d9SAndroid Build Coastguard Worker        args = mock.MagicMock()
147*800a58d9SAndroid Build Coastguard Worker        instance_obj = mock.MagicMock()
148*800a58d9SAndroid Build Coastguard Worker        # Test case with provided instance name.
149*800a58d9SAndroid Build Coastguard Worker        args.instance_name = "instance_1"
150*800a58d9SAndroid Build Coastguard Worker        args.file_name = "file1.log"
151*800a58d9SAndroid Build Coastguard Worker        args.no_prompt = True
152*800a58d9SAndroid Build Coastguard Worker        self.Patch(config, "GetAcloudConfig", return_value=cfg)
153*800a58d9SAndroid Build Coastguard Worker        self.Patch(list_instances, "GetInstancesFromInstanceNames",
154*800a58d9SAndroid Build Coastguard Worker                   return_value=[instance_obj])
155*800a58d9SAndroid Build Coastguard Worker        pull.Run(args)
156*800a58d9SAndroid Build Coastguard Worker        mock_pull_file.assert_has_calls([
157*800a58d9SAndroid Build Coastguard Worker            mock.call(cfg, instance_obj, args.file_name, args.no_prompt)])
158*800a58d9SAndroid Build Coastguard Worker
159*800a58d9SAndroid Build Coastguard Worker        # Test case for user select one instance to pull log.
160*800a58d9SAndroid Build Coastguard Worker        selected_instance = mock.MagicMock()
161*800a58d9SAndroid Build Coastguard Worker        self.Patch(list_instances, "ChooseOneRemoteInstance",
162*800a58d9SAndroid Build Coastguard Worker                   return_value=selected_instance)
163*800a58d9SAndroid Build Coastguard Worker        args.instance_name = None
164*800a58d9SAndroid Build Coastguard Worker        pull.Run(args)
165*800a58d9SAndroid Build Coastguard Worker        mock_pull_file.assert_has_calls([
166*800a58d9SAndroid Build Coastguard Worker            mock.call(cfg, selected_instance, args.file_name, args.no_prompt)])
167*800a58d9SAndroid Build Coastguard Worker
168*800a58d9SAndroid Build Coastguard Worker
169*800a58d9SAndroid Build Coastguard Workerif __name__ == '__main__':
170*800a58d9SAndroid Build Coastguard Worker    unittest.main()
171