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