1#!/usr/bin/env python
2#
3# Copyright (C) 2017 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import os
19import unittest
20
21try:
22    from unittest import mock
23except ImportError:
24    import mock
25
26try:
27    import StringIO as string_io_module
28except ImportError:
29    import io as string_io_module
30
31from host_controller.build import build_flasher
32from host_controller.tfc import command_task
33from host_controller.tfc import device_info
34from host_controller import common
35from host_controller import console
36
37
38class ConsoleTest(unittest.TestCase):
39    """A test for console.Console.
40
41    Attribute:
42        _out_file: The console output buffer.
43        _host_controller: A mock tfc_host_controller.HostController.
44        _build_provider_pab: A mock build_provider_pab.BuildProviderPAB.
45        _tfc_client: A mock tfc_client.TfcClient.
46        _vti_client A mock vti_endpoint_client.VtiEndpointClient.
47        _console: The console being tested.
48    """
49    _DEVICES = [
50        device_info.DeviceInfo(
51            device_serial="ABC001",
52            run_target="sailfish",
53            state="Available",
54            build_id="111111",
55            sdk_version="27")
56    ]
57    _TASKS = [
58        command_task.CommandTask(
59            request_id="1",
60            task_id="1-0",
61            command_id="2",
62            command_line="vts -m SampleShellTest",
63            device_serials=["ABC001"])
64    ]
65
66    def setUp(self):
67        """Creates the console."""
68        self._out_file = string_io_module.StringIO()
69        self._host_controller = mock.Mock()
70        self._build_provider_pab = mock.Mock()
71        self._tfc_client = mock.Mock()
72        self._vti_client = mock.Mock()
73        self._console = console.Console(
74            self._vti_client,
75            self._tfc_client,
76            self._build_provider_pab, [self._host_controller],
77            None,
78            out_file=self._out_file)
79        self._console.device_image_info = {}
80
81    def tearDown(self):
82        """Closes the output file."""
83        self._out_file.close()
84
85    def _IssueCommand(self, command_line):
86        """Issues a command in the console.
87
88        Args:
89            command_line: A string, the input to the console.
90
91        Returns:
92            A string, the output of the console.
93        """
94        out_position = self._out_file.tell()
95        self._console.onecmd(command_line)
96        self._out_file.seek(out_position)
97        return self._out_file.read()
98
99    def testLease(self):
100        """Tests the lease command."""
101        self._host_controller.LeaseCommandTasks.return_value = self._TASKS
102        output = self._IssueCommand("lease")
103        expected = (
104            "request_id  command_id  task_id  device_serials  command_line          \n"
105            "1           2           1-0      ABC001          vts -m SampleShellTest\n"
106        )
107        self.assertEqual(expected, output)
108        output = self._IssueCommand("lease --host 0")
109        self.assertEqual(expected, output)
110
111    def testRequest(self):
112        """Tests the request command."""
113        user = "user0"
114        cluster = "cluster0"
115        run_target = "sailfish"
116        command_line = "vts -m SampleShellTest"
117        self._IssueCommand("request --user %s --cluster %s --run-target %s "
118                           "-- %s" % (user, cluster, run_target, command_line))
119        req = self._tfc_client.NewRequest.call_args[0][0]
120        self.assertEqual(user, req.user)
121        self.assertEqual(cluster, req.cluster)
122        self.assertEqual(run_target, req.run_target)
123        self.assertEqual(command_line, req.command_line)
124
125    def testListHosts(self):
126        """Tests the list command."""
127        self._host_controller.hostname = "host0"
128        output = self._IssueCommand("list hosts")
129        self.assertEqual("index  name\n" "[  0]  host0\n", output)
130
131    def testListDevices(self):
132        """Tests the list command."""
133        self._host_controller.ListDevices.return_value = self._DEVICES
134        self._host_controller.hostname = "host0"
135        output = self._IssueCommand("list devices")
136        expected = (
137            "[  0]  host0\n"
138            "device_serial  state      run_target  build_id  sdk_version  stub\n"
139            "ABC001         Available  sailfish    111111    27               \n"
140        )
141        self.assertEqual(expected, output)
142        output = self._IssueCommand("list devices --host 0")
143        self.assertEqual(expected, output)
144
145    def testWrongHostIndex(self):
146        """Tests host index out of range."""
147        output = self._IssueCommand("list devices --host 1")
148        expected = "IndexError: "
149        self.assertTrue(output.startswith(expected))
150        output = self._IssueCommand("lease --host 1")
151        self.assertTrue(output.startswith(expected))
152
153    @mock.patch('host_controller.build.build_flasher.BuildFlasher')
154    def testFetchPOSTAndFlash(self, mock_class):
155        """Tests fetching from pab and flashing."""
156        self._build_provider_pab.GetArtifact.return_value = ({
157            "system.img":
158            "/mock/system.img",
159            "odm.img":
160            "/mock/odm.img"
161        }, {}, {
162            "build_id":
163            "build_id"
164        }, {})
165        self._build_provider_pab.GetFetchedArtifactType.return_value = common._ARTIFACT_TYPE_DEVICE
166        self._IssueCommand(
167            "fetch --branch=aosp-master-ndk --target=darwin_mac "
168            "--account_id=100621237 "
169            "--artifact_name=foo-{build_id}.tar.bz2 --method=POST")
170        self._build_provider_pab.GetArtifact.assert_called_with(
171            account_id='100621237',
172            branch='aosp-master-ndk',
173            target='darwin_mac',
174            artifact_name='foo-{build_id}.tar.bz2',
175            build_id='latest',
176            method='POST',
177            full_device_images=False)
178        self.assertEqual(self._console.device_image_info, {
179            "system.img": "/mock/system.img",
180            "odm.img": "/mock/odm.img"
181        })
182
183        flasher = mock.Mock()
184        mock_class.return_value = flasher
185        self._IssueCommand("flash --current system=system.img odm=odm.img")
186        flasher.Flash.assert_called_with({
187            "system": "/mock/system.img",
188            "odm": "/mock/odm.img"
189        }, False)
190
191    def testFetchAndEnvironment(self):
192        """Tests fetching from pab and check stored os environment"""
193        build_id_return = "4328532"
194        target_return = "darwin_mac"
195        expected_fetch_info = {"build_id": build_id_return}
196
197        self._build_provider_pab.GetArtifact.return_value = ({
198            "system.img":
199            "/mock/system.img",
200            "odm.img":
201            "/mock/odm.img"
202        }, {}, expected_fetch_info, {})
203        self._IssueCommand(
204            "fetch --branch=aosp-master-ndk --target=%s "
205            "--account_id=100621237 "
206            "--artifact_name=foo-{id}.tar.bz2 --method=POST" % target_return)
207        self._build_provider_pab.GetArtifact.assert_called_with(
208            account_id='100621237',
209            branch='aosp-master-ndk',
210            target='darwin_mac',
211            artifact_name='foo-{id}.tar.bz2',
212            build_id='latest',
213            full_device_images=False,
214            method='POST')
215
216        expected = expected_fetch_info["build_id"]
217        self.assertEqual(build_id_return, expected)
218
219    @mock.patch('host_controller.build.build_flasher.BuildFlasher')
220    def testFlashGSI(self, mock_class):
221        flasher = mock.Mock()
222        mock_class.return_value = flasher
223        self._IssueCommand("flash --gsi=system.img")
224        flasher.FlashGSI.assert_called_with(
225            'system.img', None, skip_vbmeta=False)
226
227    @mock.patch('host_controller.build.build_flasher.BuildFlasher')
228    def testFlashGSIWithVbmeta(self, mock_class):
229        flasher = mock.Mock()
230        mock_class.return_value = flasher
231        self._IssueCommand("flash --gsi=system.img --vbmeta=vbmeta.img")
232        flasher.FlashGSI.assert_called_with(
233            'system.img', 'vbmeta.img', skip_vbmeta=False)
234
235    @mock.patch('host_controller.build.build_flasher.BuildFlasher')
236    def testFlashall(self, mock_class):
237        flasher = mock.Mock()
238        mock_class.return_value = flasher
239        self._IssueCommand("flash --build_dir=path/to/dir/")
240        flasher.Flashall.assert_called_with('path/to/dir/')
241
242    @mock.patch('host_controller.command_processor.command_flash.importlib')
243    @mock.patch('host_controller.command_processor.command_flash.issubclass')
244    def testImportFlasher(self, mock_issubclass, mock_importlib):
245        mock_issubclass.return_value = True
246        flasher_module = mock.Mock()
247        flasher = mock.Mock()
248        mock_importlib.import_module.return_value = flasher_module
249        flasher_module.Flasher.return_value = flasher
250        self._IssueCommand("flash --serial ABC001 "
251                           "--flasher_type test.flasher.Flasher "
252                           "--flasher_path /test/flasher "
253                           "-- --unit test")
254        mock_issubclass.assert_called_once_with(flasher_module.Flasher,
255                                                build_flasher.BuildFlasher)
256        mock_importlib.import_module.assert_called_with("test.flasher")
257        flasher_module.Flasher.assert_called_with("ABC001", "/test/flasher")
258        flasher.Flash.assert_called_with({}, {}, "--unit", "test")
259        flasher.WaitForDevice.assert_called_with()
260
261
262if __name__ == "__main__":
263    unittest.main()
264