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 18 import os 19 import unittest 20 21 try: 22 from unittest import mock 23 except ImportError: 24 import mock 25 26 try: 27 import StringIO as string_io_module 28 except ImportError: 29 import io as string_io_module 30 31 from host_controller.build import build_flasher 32 from host_controller.tfc import command_task 33 from host_controller.tfc import device_info 34 from host_controller import console 35 36 37 class ConsoleTest(unittest.TestCase): 38 """A test for console.Console. 39 40 Attribute: 41 _out_file: The console output buffer. 42 _host_controller: A mock tfc_host_controller.HostController. 43 _build_provider_pab: A mock build_provider_pab.BuildProviderPAB. 44 _tfc_client: A mock tfc_client.TfcClient. 45 _vti_client A mock vti_endpoint_client.VtiEndpointClient. 46 _console: The console being tested. 47 """ 48 _DEVICES = [ 49 device_info.DeviceInfo( 50 device_serial="ABC001", 51 run_target="sailfish", 52 state="Available", 53 build_id="111111", 54 sdk_version="27") 55 ] 56 _TASKS = [ 57 command_task.CommandTask( 58 request_id="1", 59 task_id="1-0", 60 command_id="2", 61 command_line="vts -m SampleShellTest", 62 device_serials=["ABC001"]) 63 ] 64 65 def setUp(self): 66 """Creates the console.""" 67 self._out_file = string_io_module.StringIO() 68 self._host_controller = mock.Mock() 69 self._build_provider_pab = mock.Mock() 70 self._tfc_client = mock.Mock() 71 self._vti_client = mock.Mock() 72 self._console = console.Console( 73 self._vti_client, self._tfc_client, self._build_provider_pab, 74 [self._host_controller], 75 None, self._out_file) 76 77 def tearDown(self): 78 """Closes the output file.""" 79 self._out_file.close() 80 81 def _IssueCommand(self, command_line): 82 """Issues a command in the console. 83 84 Args: 85 command_line: A string, the input to the console. 86 87 Returns: 88 A string, the output of the console. 89 """ 90 out_position = self._out_file.tell() 91 self._console.onecmd(command_line) 92 self._out_file.seek(out_position) 93 return self._out_file.read() 94 95 def testLease(self): 96 """Tests the lease command.""" 97 self._host_controller.LeaseCommandTasks.return_value = self._TASKS 98 output = self._IssueCommand("lease") 99 expected = ( 100 "request_id command_id task_id device_serials command_line \n" 101 "1 2 1-0 ABC001 vts -m SampleShellTest\n" 102 ) 103 self.assertEqual(expected, output) 104 output = self._IssueCommand("lease --host 0") 105 self.assertEqual(expected, output) 106 107 def testRequest(self): 108 """Tests the request command.""" 109 user = "user0" 110 cluster = "cluster0" 111 run_target = "sailfish" 112 command_line = "vts -m SampleShellTest" 113 self._IssueCommand("request --user %s --cluster %s --run-target %s " 114 "-- %s" % (user, cluster, run_target, command_line)) 115 req = self._tfc_client.NewRequest.call_args[0][0] 116 self.assertEqual(user, req.user) 117 self.assertEqual(cluster, req.cluster) 118 self.assertEqual(run_target, req.run_target) 119 self.assertEqual(command_line, req.command_line) 120 121 def testListHosts(self): 122 """Tests the list command.""" 123 self._host_controller.hostname = "host0" 124 output = self._IssueCommand("list hosts") 125 self.assertEqual("index name\n" "[ 0] host0\n", output) 126 127 def testListDevices(self): 128 """Tests the list command.""" 129 self._host_controller.ListDevices.return_value = self._DEVICES 130 self._host_controller.hostname = "host0" 131 output = self._IssueCommand("list devices") 132 expected = ( 133 "[ 0] host0\n" 134 "device_serial state run_target build_id sdk_version stub\n" 135 "ABC001 Available sailfish 111111 27 \n" 136 ) 137 self.assertEqual(expected, output) 138 output = self._IssueCommand("list devices --host 0") 139 self.assertEqual(expected, output) 140 141 def testWrongHostIndex(self): 142 """Tests host index out of range.""" 143 output = self._IssueCommand("list devices --host 1") 144 expected = "IndexError: " 145 self.assertTrue(output.startswith(expected)) 146 output = self._IssueCommand("lease --host 1") 147 self.assertTrue(output.startswith(expected)) 148 149 @mock.patch('host_controller.console.build_flasher.BuildFlasher') 150 def testFetchPOSTAndFlash(self, mock_class): 151 """Tests fetching from pab and flashing.""" 152 self._build_provider_pab.GetArtifact.return_value = ({ 153 "system.img": 154 "/mock/system.img", 155 "odm.img": 156 "/mock/odm.img" 157 }, {}, {"build_id": "build_id"}, {}) 158 self._IssueCommand( 159 "fetch --branch=aosp-master-ndk --target=darwin_mac " 160 "--account_id=100621237 " 161 "--artifact_name=foo-{id}.tar.bz2 --method=POST") 162 self._build_provider_pab.GetArtifact.assert_called_with( 163 account_id='100621237', 164 branch='aosp-master-ndk', 165 target='darwin_mac', 166 artifact_name='foo-{id}.tar.bz2', 167 build_id='latest', 168 method='POST') 169 170 flasher = mock.Mock() 171 mock_class.return_value = flasher 172 self._IssueCommand("flash --current system=system.img odm=odm.img") 173 flasher.Flash.assert_called_with({ 174 "system": "/mock/system.img", 175 "odm": "/mock/odm.img" 176 }) 177 178 def testFetchAndEnvironment(self): 179 """Tests fetching from pab and check stored os environment""" 180 build_id_return = "4328532" 181 target_return = "darwin_mac" 182 expected_fetch_info = {"build_id": build_id_return} 183 184 self._build_provider_pab.GetArtifact.return_value = ({ 185 "system.img": 186 "/mock/system.img", 187 "odm.img": 188 "/mock/odm.img" 189 }, {}, expected_fetch_info, {}) 190 self._IssueCommand("fetch --branch=aosp-master-ndk --target=%s " 191 "--account_id=100621237 " 192 "--artifact_name=foo-{id}.tar.bz2 --method=POST" 193 % target_return) 194 self._build_provider_pab.GetArtifact.assert_called_with( 195 account_id='100621237', 196 branch='aosp-master-ndk', 197 target='darwin_mac', 198 artifact_name='foo-{id}.tar.bz2', 199 build_id='latest', 200 method='POST') 201 202 expected = expected_fetch_info["build_id"] 203 self.assertEqual(build_id_return, expected) 204 205 @mock.patch('host_controller.console.build_flasher.BuildFlasher') 206 def testFlashGSI(self, mock_class): 207 flasher = mock.Mock() 208 mock_class.return_value = flasher 209 self._IssueCommand("flash --gsi=system.img") 210 flasher.FlashGSI.assert_called_with('system.img', None) 211 212 @mock.patch('host_controller.console.build_flasher.BuildFlasher') 213 def testFlashGSIWithVbmeta(self, mock_class): 214 flasher = mock.Mock() 215 mock_class.return_value = flasher 216 self._IssueCommand("flash --gsi=system.img --vbmeta=vbmeta.img") 217 flasher.FlashGSI.assert_called_with('system.img', 'vbmeta.img') 218 219 @mock.patch('host_controller.console.build_flasher.BuildFlasher') 220 def testFlashall(self, mock_class): 221 flasher = mock.Mock() 222 mock_class.return_value = flasher 223 self._IssueCommand("flash --build_dir=path/to/dir/") 224 flasher.Flashall.assert_called_with('path/to/dir/') 225 226 @mock.patch('host_controller.console.importlib') 227 @mock.patch('host_controller.console.issubclass') 228 def testImportFlasher(self, mock_issubclass, mock_importlib): 229 mock_issubclass.return_value = True 230 flasher_module = mock.Mock() 231 flasher = mock.Mock() 232 mock_importlib.import_module.return_value = flasher_module 233 flasher_module.Flasher.return_value = flasher 234 self._IssueCommand("flash --serial ABC001 " 235 "--flasher_type test.flasher.Flasher " 236 "--flasher_path /test/flasher " 237 "-- --unit test") 238 mock_issubclass.assert_called_once_with( 239 flasher_module.Flasher, build_flasher.BuildFlasher) 240 mock_importlib.import_module.assert_called_with("test.flasher") 241 flasher_module.Flasher.assert_called_with("ABC001", "/test/flasher") 242 flasher.Flash.assert_called_with({}, {}, "--unit", "test") 243 flasher.WaitForDevice.assert_called_with() 244 245 246 if __name__ == "__main__": 247 unittest.main() 248