1# Copyright 2017 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"""Unit test for fastboot interface using sh library."""
16import unittest
17
18import fastboot_exceptions
19import fastbootsh
20from mock import patch
21import sh
22
23
24class FastbootShTest(unittest.TestCase):
25
26  ATFA_TEST_SERIAL = 'ATFA_TEST_SERIAL'
27  TEST_MESSAGE_FAILURE = 'FAIL: TEST MESSAGE'
28  TEST_MESSAGE_SUCCESS = 'OKAY: TEST MESSAGE'
29  TEST_SERIAL = 'TEST_SERIAL'
30
31  TEST_VAR = 'VAR1'
32
33  class TestError(sh.ErrorReturnCode):
34
35    def __init__(self):
36      pass
37
38  def setUp(self):
39    pass
40
41  # Test FastbootDevice.ListDevices
42  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
43  def testListDevicesOneDevice(self, mock_fastboot_commands):
44    mock_fastboot_commands.return_value = self.TEST_SERIAL + '\tfastboot'
45    device_serial_numbers = fastbootsh.FastbootDevice.ListDevices()
46    mock_fastboot_commands.assert_called_once_with('devices')
47    self.assertEqual(1, len(device_serial_numbers))
48    self.assertEqual(self.TEST_SERIAL, device_serial_numbers[0])
49
50  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
51  def testListDevicesTwoDevices(self, mock_fastboot_commands):
52    mock_fastboot_commands.return_value = (self.TEST_SERIAL + '\tfastboot\n' +
53                                           self.ATFA_TEST_SERIAL + '\tfastboot')
54    device_serial_numbers = fastbootsh.FastbootDevice.ListDevices()
55    mock_fastboot_commands.assert_called_once_with('devices')
56    self.assertEqual(2, len(device_serial_numbers))
57    self.assertEqual(self.TEST_SERIAL, device_serial_numbers[0])
58    self.assertEqual(self.ATFA_TEST_SERIAL, device_serial_numbers[1])
59
60  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
61  def testListDevicesMultiDevices(self, mock_fastboot_commands):
62    one_device = self.TEST_SERIAL + '\tfastboot'
63    result = one_device
64    for _ in range(0, 9):
65      result += '\n' + one_device
66    mock_fastboot_commands.return_value = result
67    device_serial_numbers = fastbootsh.FastbootDevice.ListDevices()
68    mock_fastboot_commands.assert_called_once_with('devices')
69    self.assertEqual(10, len(device_serial_numbers))
70
71  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
72  def testListDevicesNone(self, mock_fastboot_commands):
73    mock_fastboot_commands.return_value = ''
74    device_serial_numbers = fastbootsh.FastbootDevice.ListDevices()
75    mock_fastboot_commands.assert_called_once_with('devices')
76    self.assertEqual(0, len(device_serial_numbers))
77
78  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
79  def testListDevicesFailure(self, mock_fastboot_commands):
80    mock_error = self.TestError()
81    mock_error.stderr = self.TEST_MESSAGE_FAILURE
82    mock_fastboot_commands.side_effect = mock_error
83    with self.assertRaises(fastboot_exceptions.FastbootFailure) as e:
84      fastbootsh.FastbootDevice.ListDevices()
85    self.assertEqual(self.TEST_MESSAGE_FAILURE, str(e.exception))
86
87  # Test FastbootDevice.Oem
88  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
89  def testOem(self, mock_fastboot_commands):
90    mock_fastboot_commands.return_value = self.TEST_MESSAGE_SUCCESS
91    command = 'TEST COMMAND'
92    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
93    message = device.Oem(command, False)
94    mock_fastboot_commands.assert_called_once_with('-s',
95                                                   self.TEST_SERIAL,
96                                                   'oem',
97                                                   command,
98                                                   _err_to_out=False)
99    self.assertEqual(self.TEST_MESSAGE_SUCCESS, message)
100
101  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
102  def testOemErrToOut(self, mock_fastboot_commands):
103    mock_fastboot_commands.return_value = self.TEST_MESSAGE_SUCCESS
104    command = 'TEST COMMAND'
105    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
106    message = device.Oem(command, True)
107    mock_fastboot_commands.assert_called_once_with('-s',
108                                                   self.TEST_SERIAL,
109                                                   'oem',
110                                                   command,
111                                                   _err_to_out=True)
112    self.assertEqual(self.TEST_MESSAGE_SUCCESS, message)
113
114  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
115  def testOemFailure(self, mock_fastboot_commands):
116    mock_error = self.TestError()
117    mock_error.stderr = self.TEST_MESSAGE_FAILURE
118    mock_fastboot_commands.side_effect = mock_error
119    command = 'TEST COMMAND'
120    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
121    with self.assertRaises(fastboot_exceptions.FastbootFailure) as e:
122      device.Oem(command, False)
123    self.assertEqual(self.TEST_MESSAGE_FAILURE, str(e.exception))
124
125  # Test FastbootDevice.Upload
126  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
127  def testUpload(self, mock_fastboot_commands):
128    mock_fastboot_commands.return_value = self.TEST_MESSAGE_SUCCESS
129    command = 'TEST COMMAND'
130    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
131    message = device.Upload(command)
132    mock_fastboot_commands.assert_called_once_with('-s',
133                                                   self.TEST_SERIAL,
134                                                   'get_staged',
135                                                   command)
136    self.assertEqual(self.TEST_MESSAGE_SUCCESS, message)
137
138  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
139  def testUploadFailure(self, mock_fastboot_commands):
140    mock_error = self.TestError()
141    mock_error.stderr = self.TEST_MESSAGE_FAILURE
142    mock_fastboot_commands.side_effect = mock_error
143    command = 'TEST COMMAND'
144    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
145    with self.assertRaises(fastboot_exceptions.FastbootFailure) as e:
146      device.Upload(command)
147    self.assertEqual(self.TEST_MESSAGE_FAILURE, str(e.exception))
148
149  # Test FastbootDevice.Download
150  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
151  def testDownload(self, mock_fastboot_commands):
152    mock_fastboot_commands.return_value = self.TEST_MESSAGE_SUCCESS
153    command = 'TEST COMMAND'
154    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
155    message = device.Download(command)
156    mock_fastboot_commands.assert_called_once_with('-s',
157                                                   self.TEST_SERIAL,
158                                                   'stage',
159                                                   command)
160    self.assertEqual(self.TEST_MESSAGE_SUCCESS, message)
161
162  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
163  def testDownloadFailure(self, mock_fastboot_commands):
164    mock_error = self.TestError()
165    mock_error.stderr = self.TEST_MESSAGE_FAILURE
166    mock_fastboot_commands.side_effect = mock_error
167    command = 'TEST COMMAND'
168    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
169    with self.assertRaises(fastboot_exceptions.FastbootFailure) as e:
170      device.Download(command)
171    self.assertEqual(self.TEST_MESSAGE_FAILURE, str(e.exception))
172
173  # Test FastbootDevice.GetVar
174  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
175  def testGetVar(self, mock_fastboot_commands):
176    mock_fastboot_commands.return_value = self.TEST_VAR + ': ' + 'abcd'
177    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
178    message = device.GetVar(self.TEST_VAR)
179    mock_fastboot_commands.assert_called_once_with('-s',
180                                                   self.TEST_SERIAL,
181                                                   'getvar',
182                                                   self.TEST_VAR,
183                                                   _err_to_out=True)
184    self.assertEqual('abcd', message)
185
186  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
187  def testGetVarFailure(self, mock_fastboot_commands):
188    mock_error = self.TestError()
189    mock_error.stdout = self.TEST_MESSAGE_FAILURE
190    mock_fastboot_commands.side_effect = mock_error
191    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
192    with self.assertRaises(fastboot_exceptions.FastbootFailure) as e:
193      device.GetVar(self.TEST_VAR)
194    self.assertEqual(self.TEST_MESSAGE_FAILURE, str(e.exception))
195
196  # Test FastbootDevice.Reboot
197  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
198  def testReboot(self, mock_fastboot_commands):
199    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
200    device.Reboot()
201    mock_fastboot_commands.assert_called_once_with('-s',
202                                                   self.TEST_SERIAL,
203                                                   'reboot-bootloader')
204
205  @patch('fastbootsh.FastbootDevice.fastboot_command', create=True)
206  def testRebootFailure(self, mock_fastboot_commands):
207    mock_error = self.TestError()
208    mock_error.stderr = self.TEST_MESSAGE_FAILURE
209    mock_fastboot_commands.side_effect = mock_error
210    device = fastbootsh.FastbootDevice(self.TEST_SERIAL)
211    with self.assertRaises(fastboot_exceptions.FastbootFailure) as e:
212      device.Reboot()
213    self.assertEqual(self.TEST_MESSAGE_FAILURE, str(e.exception))
214
215if __name__ == '__main__':
216  unittest.main()
217