1#!/usr/bin/env python
2# Copyright 2016 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6# pylint: disable=protected-access
7
8"""
9Unit tests for the contents of find_usb_devices.py.
10
11Device tree for these tests is as follows:
12Bus 001:
131: Device 011 "foo"
142: Device 012 "bar"
153: Device 013 "baz"
16
17Bus 002:
181: Device 011 "quux"
192: Device 020 "My Test HUB" #hub 1
202:1: Device 021 "battor_p7_h1_t0" #physical port 7 on hub 1, on ttyUSB0
212:3: Device 022 "battor_p5_h1_t1" #physical port 5 on hub 1, on ttyUSB1
222:4: Device 023 "My Test Internal HUB" #internal section of hub 1
232:4:2: Device 024 "battor_p3_h1_t2" #physical port 3 on hub 1, on ttyUSB2
242:4:3: Device 026 "Not a Battery Monitor" #physical port 1 on hub 1, on ttyUSB3
252:4:4: Device 025 "battor_p1_h1_t3" #physical port 1 on hub 1, on ttyUSB3
263: Device 100 "My Test HUB" #hub 2
273:4: Device 101 "My Test Internal HUB" #internal section of hub 2
283:4:4: Device 102 "battor_p1_h2_t4" #physical port 1 on hub 2, on ttyusb4
29"""
30
31import logging
32import os
33import unittest
34
35from devil import devil_env
36from devil.utils import battor_device_mapping
37from devil.utils import find_usb_devices
38from devil.utils import lsusb
39from devil.utils import usb_hubs
40
41with devil_env.SysPath(devil_env.PYMOCK_PATH):
42  import mock # pylint: disable=import-error
43
44# Output of lsusb.lsusb().
45# We just test that the dictionary is working by creating an
46# "ID number" equal to (bus_num*1000)+device_num and seeing if
47# it is picked up correctly. Also we test the description
48
49DEVLIST = [(1, 11, 'foo'),
50           (1, 12, 'bar'),
51           (1, 13, 'baz'),
52           (2, 11, 'quux'),
53           (2, 20, 'My Test HUB'),
54           (2, 21, 'ID 0403:6001 battor_p7_h1_t0'),
55           (2, 22, 'ID 0403:6001 battor_p5_h1_t1'),
56           (2, 23, 'My Test Internal HUB'),
57           (2, 24, 'ID 0403:6001 battor_p3_h1_t2'),
58           (2, 25, 'ID 0403:6001 battor_p1_h1_t3'),
59           (2, 26, 'Not a Battery Monitor'),
60           (2, 100, 'My Test HUB'),
61           (2, 101, 'My Test Internal HUB'),
62           (2, 102, 'ID 0403:6001 battor_p1_h1_t4')]
63
64LSUSB_OUTPUT = [
65  {'bus': b, 'device': d, 'desc': t, 'id': (1000*b)+d}
66       for (b, d, t) in DEVLIST]
67
68
69# Note: "Lev", "Cnt", "Spd", and "MxCh" are not used by parser,
70# so we just leave them as zeros here. Also note that the port
71# numbers reported here start at 0, so they're 1 less than the
72# port numbers reported elsewhere.
73USB_DEVICES_OUTPUT = '''
74T:  Bus=01 Lev=00 Prnt=00 Port=00 Cnt=00 Dev#= 11 Spd=000 MxCh=00
75S:  SerialNumber=FooSerial
76T:  Bus=01 Lev=00 Prnt=00 Port=01 Cnt=00 Dev#= 12 Spd=000 MxCh=00
77S:  SerialNumber=BarSerial
78T:  Bus=01 Lev=00 Prnt=00 Port=02 Cnt=00 Dev#= 13 Spd=000 MxCh=00
79S:  SerialNumber=BazSerial
80
81T:  Bus=02 Lev=00 Prnt=00 Port=00 Cnt=00 Dev#= 11 Spd=000 MxCh=00
82
83T:  Bus=02 Lev=00 Prnt=00 Port=01 Cnt=00 Dev#= 20 Spd=000 MxCh=00
84T:  Bus=02 Lev=00 Prnt=20 Port=00 Cnt=00 Dev#= 21 Spd=000 MxCh=00
85S:  SerialNumber=Battor0
86T:  Bus=02 Lev=00 Prnt=20 Port=02 Cnt=00 Dev#= 22 Spd=000 MxCh=00
87S:  SerialNumber=Battor1
88T:  Bus=02 Lev=00 Prnt=20 Port=03 Cnt=00 Dev#= 23 Spd=000 MxCh=00
89T:  Bus=02 Lev=00 Prnt=23 Port=01 Cnt=00 Dev#= 24 Spd=000 MxCh=00
90S:  SerialNumber=Battor2
91T:  Bus=02 Lev=00 Prnt=23 Port=03 Cnt=00 Dev#= 25 Spd=000 MxCh=00
92S:  SerialNumber=Battor3
93T:  Bus=02 Lev=00 Prnt=23 Port=02 Cnt=00 Dev#= 26 Spd=000 MxCh=00
94
95T:  Bus=02 Lev=00 Prnt=00 Port=02 Cnt=00 Dev#=100 Spd=000 MxCh=00
96T:  Bus=02 Lev=00 Prnt=100 Port=03 Cnt=00 Dev#=101 Spd=000 MxCh=00
97T:  Bus=02 Lev=00 Prnt=101 Port=03 Cnt=00 Dev#=102 Spd=000 MxCh=00
98'''
99
100RAW_LSUSB_OUTPUT = '''
101Bus 001 Device 011: FAST foo
102Bus 001 Device 012: FAST bar
103Bus 001 Device 013: baz
104Bus 002 Device 011: quux
105Bus 002 Device 020: My Test HUB
106Bus 002 Device 021: ID 0403:6001 battor_p7_h1_t0
107Bus 002 Device 022: ID 0403:6001 battor_p5_h1_t1
108Bus 002 Device 023: My Test Internal HUB
109Bus 002 Device 024: ID 0403:6001 battor_p3_h1_t2
110Bus 002 Device 025: ID 0403:6001 battor_p1_h1_t3
111Bus 002 Device 026: Not a Battery Monitor
112Bus 002 Device 100: My Test HUB
113Bus 002 Device 101: My Test Internal HUB
114Bus 002 Device 102: ID 0403:6001 battor_p1_h1_t4
115'''
116
117LIST_TTY_OUTPUT = '''
118ttyUSB0
119Something-else-0
120ttyUSB1
121ttyUSB2
122Something-else-1
123ttyUSB3
124ttyUSB4
125Something-else-2
126ttyUSB5
127'''
128
129# Note: The real output will have multiple lines with
130# ATTRS{busnum} and ATTRS{devnum}, but only the first
131# one counts. Thus the test output duplicates this.
132UDEVADM_USBTTY0_OUTPUT = '''
133ATTRS{busnum}=="2"
134ATTRS{devnum}=="21"
135ATTRS{busnum}=="0"
136ATTRS{devnum}=="0"
137'''
138
139UDEVADM_USBTTY1_OUTPUT = '''
140ATTRS{busnum}=="2"
141ATTRS{devnum}=="22"
142ATTRS{busnum}=="0"
143ATTRS{devnum}=="0"
144'''
145
146UDEVADM_USBTTY2_OUTPUT = '''
147ATTRS{busnum}=="2"
148ATTRS{devnum}=="24"
149ATTRS{busnum}=="0"
150ATTRS{devnum}=="0"
151'''
152
153UDEVADM_USBTTY3_OUTPUT = '''
154ATTRS{busnum}=="2"
155ATTRS{devnum}=="25"
156ATTRS{busnum}=="0"
157ATTRS{devnum}=="0"
158'''
159
160UDEVADM_USBTTY4_OUTPUT = '''
161ATTRS{busnum}=="2"
162ATTRS{devnum}=="102"
163ATTRS{busnum}=="0"
164ATTRS{devnum}=="0"
165'''
166
167UDEVADM_USBTTY5_OUTPUT = '''
168ATTRS{busnum}=="2"
169ATTRS{devnum}=="26"
170ATTRS{busnum}=="0"
171ATTRS{devnum}=="0"
172'''
173
174UDEVADM_OUTPUT_DICT = {
175  'ttyUSB0': UDEVADM_USBTTY0_OUTPUT,
176  'ttyUSB1': UDEVADM_USBTTY1_OUTPUT,
177  'ttyUSB2': UDEVADM_USBTTY2_OUTPUT,
178  'ttyUSB3': UDEVADM_USBTTY3_OUTPUT,
179  'ttyUSB4': UDEVADM_USBTTY4_OUTPUT,
180  'ttyUSB5': UDEVADM_USBTTY5_OUTPUT}
181
182# Identification criteria for Plugable 7-Port Hub
183def isTestHub(node):
184  """Check if a node is a Plugable 7-Port Hub
185  (Model USB2-HUB7BC)
186  The topology of this device is a 4-port hub,
187  with another 4-port hub connected on port 4.
188  """
189  if not isinstance(node, find_usb_devices.USBDeviceNode):
190    return False
191  if 'Test HUB' not in node.desc:
192    return False
193  if not node.HasPort(4):
194    return False
195  return 'Test Internal HUB' in node.PortToDevice(4).desc
196
197TEST_HUB = usb_hubs.HubType(isTestHub,
198                            {1:7,
199                             2:6,
200                             3:5,
201                             4:{1:4, 2:3, 3:2, 4:1}})
202
203class USBScriptTest(unittest.TestCase):
204  def setUp(self):
205    find_usb_devices._GetTtyUSBInfo = mock.Mock(
206        side_effect=lambda x: UDEVADM_OUTPUT_DICT[x])
207    find_usb_devices._GetParsedLSUSBOutput = mock.Mock(
208        return_value=LSUSB_OUTPUT)
209    find_usb_devices._GetUSBDevicesOutput = mock.Mock(
210        return_value=USB_DEVICES_OUTPUT)
211    find_usb_devices._GetCommList = mock.Mock(
212        return_value=LIST_TTY_OUTPUT)
213    lsusb.raw_lsusb = mock.Mock(
214        return_value=RAW_LSUSB_OUTPUT)
215
216  def testIsBattor(self):
217    bd = find_usb_devices.GetBusNumberToDeviceTreeMap()
218    self.assertTrue(battor_device_mapping.IsBattor('ttyUSB3', bd))
219    self.assertFalse(battor_device_mapping.IsBattor('ttyUSB5', bd))
220
221  def testGetBattors(self):
222    bd = find_usb_devices.GetBusNumberToDeviceTreeMap()
223    self.assertEquals(battor_device_mapping.GetBattorList(bd),
224                          ['ttyUSB0', 'ttyUSB1', 'ttyUSB2',
225                           'ttyUSB3', 'ttyUSB4'])
226
227  def testGetTTYDevices(self):
228    pp = find_usb_devices.GetAllPhysicalPortToTTYMaps([TEST_HUB])
229    result = list(pp)
230    self.assertEquals(result[0], {7:'ttyUSB0',
231                                  5:'ttyUSB1',
232                                  3:'ttyUSB2',
233                                  2:'ttyUSB5',
234                                  1:'ttyUSB3'})
235    self.assertEquals(result[1], {1:'ttyUSB4'})
236
237  def testGetPortDeviceMapping(self):
238    pp = find_usb_devices.GetAllPhysicalPortToBusDeviceMaps([TEST_HUB])
239    result = list(pp)
240    self.assertEquals(result[0], {7:(2, 21),
241                                  5:(2, 22),
242                                  3:(2, 24),
243                                  2:(2, 26),
244                                  1:(2, 25)})
245    self.assertEquals(result[1], {1:(2, 102)})
246
247  def testGetSerialMapping(self):
248    pp = find_usb_devices.GetAllPhysicalPortToSerialMaps([TEST_HUB])
249    result = list(pp)
250    self.assertEquals(result[0], {7:'Battor0',
251                                  5:'Battor1',
252                                  3:'Battor2',
253                                  1:'Battor3'})
254    self.assertEquals(result[1], {})
255
256  def testFastDeviceDescriptions(self):
257    bd = find_usb_devices.GetBusNumberToDeviceTreeMap()
258    dev_foo = bd[1].FindDeviceNumber(11)
259    dev_bar = bd[1].FindDeviceNumber(12)
260    dev_battor_p7_h1_t0 = bd[2].FindDeviceNumber(21)
261    self.assertEquals(dev_foo.desc, 'FAST foo')
262    self.assertEquals(dev_bar.desc, 'FAST bar')
263    self.assertEquals(dev_battor_p7_h1_t0.desc,
264        'ID 0403:6001 battor_p7_h1_t0')
265
266  def testDeviceDescriptions(self):
267    bd = find_usb_devices.GetBusNumberToDeviceTreeMap(fast=False)
268    dev_foo = bd[1].FindDeviceNumber(11)
269    dev_bar = bd[1].FindDeviceNumber(12)
270    dev_battor_p7_h1_t0 = bd[2].FindDeviceNumber(21)
271    self.assertEquals(dev_foo.desc, 'foo')
272    self.assertEquals(dev_bar.desc, 'bar')
273    self.assertEquals(dev_battor_p7_h1_t0.desc,
274        'ID 0403:6001 battor_p7_h1_t0')
275
276  def testDeviceInformation(self):
277    bd = find_usb_devices.GetBusNumberToDeviceTreeMap(fast=False)
278    dev_foo = bd[1].FindDeviceNumber(11)
279    dev_bar = bd[1].FindDeviceNumber(12)
280    dev_battor_p7_h1_t0 = bd[2].FindDeviceNumber(21)
281    self.assertEquals(dev_foo.info['id'], 1011)
282    self.assertEquals(dev_bar.info['id'], 1012)
283    self.assertEquals(dev_battor_p7_h1_t0.info['id'], 2021)
284
285  def testSerialNumber(self):
286    bd = find_usb_devices.GetBusNumberToDeviceTreeMap(fast=False)
287    dev_foo = bd[1].FindDeviceNumber(11)
288    dev_bar = bd[1].FindDeviceNumber(12)
289    dev_battor_p7_h1_t0 = bd[2].FindDeviceNumber(21)
290    self.assertEquals(dev_foo.serial, 'FooSerial')
291    self.assertEquals(dev_bar.serial, 'BarSerial')
292    self.assertEquals(dev_battor_p7_h1_t0.serial, 'Battor0')
293
294  def testBattorDictMapping(self):
295    map_dict = {'Phone1':'Battor1', 'Phone2':'Battor2', 'Phone3':'Battor3'}
296    a1 = battor_device_mapping.GetBattorPathFromPhoneSerial(
297             'Phone1', serial_map=map_dict)
298    a2 = battor_device_mapping.GetBattorPathFromPhoneSerial(
299             'Phone2', serial_map=map_dict)
300    a3 = battor_device_mapping.GetBattorPathFromPhoneSerial(
301             'Phone3', serial_map=map_dict)
302    self.assertEquals(a1, '/dev/ttyUSB1')
303    self.assertEquals(a2, '/dev/ttyUSB2')
304    self.assertEquals(a3, '/dev/ttyUSB3')
305
306  def testBattorDictFromFileMapping(self):
307    try:
308      map_dict = {'Phone1':'Battor1', 'Phone2':'Battor2', 'Phone3':'Battor3'}
309      curr_dir = os.path.dirname(os.path.realpath(__file__))
310      filename = os.path.join(curr_dir, 'test', 'data', 'test_write_map.json')
311      battor_device_mapping.WriteSerialMapFile(filename, map_dict)
312      a1 = battor_device_mapping.GetBattorPathFromPhoneSerial(
313               'Phone1', serial_map_file=filename)
314      a2 = battor_device_mapping.GetBattorPathFromPhoneSerial(
315               'Phone2', serial_map_file=filename)
316      a3 = battor_device_mapping.GetBattorPathFromPhoneSerial(
317               'Phone3', serial_map_file=filename)
318    finally:
319      os.remove(filename)
320    self.assertEquals(a1, '/dev/ttyUSB1')
321    self.assertEquals(a2, '/dev/ttyUSB2')
322    self.assertEquals(a3, '/dev/ttyUSB3')
323
324  def testReadSerialMapFile(self):
325    curr_dir = os.path.dirname(os.path.realpath(__file__))
326    map_dict = battor_device_mapping.ReadSerialMapFile(
327        os.path.join(curr_dir, 'test', 'data', 'test_serial_map.json'))
328    self.assertEquals(len(map_dict.keys()), 3)
329    self.assertEquals(map_dict['Phone1'], 'Battor1')
330    self.assertEquals(map_dict['Phone2'], 'Battor2')
331    self.assertEquals(map_dict['Phone3'], 'Battor3')
332
333original_PPTSM = find_usb_devices.GetAllPhysicalPortToSerialMaps
334original_PPTTM = find_usb_devices.GetAllPhysicalPortToTTYMaps
335original_GBL = battor_device_mapping.GetBattorList
336original_GBNDM = find_usb_devices.GetBusNumberToDeviceTreeMap
337original_IB = battor_device_mapping.IsBattor
338original_GBSM = battor_device_mapping.GetBattorSerialNumbers
339
340def setup_battor_test(serial, tty, battor, bser=None):
341  serial_mapper = mock.Mock(return_value=serial)
342  tty_mapper = mock.Mock(return_value=tty)
343  battor_lister = mock.Mock(return_value=battor)
344  devtree = mock.Mock(return_value=None)
345  is_battor = mock.Mock(side_effect=lambda x, y: x in battor)
346  battor_serials = mock.Mock(return_value=bser)
347  find_usb_devices.GetAllPhysicalPortToSerialMaps = serial_mapper
348  find_usb_devices.GetAllPhysicalPortToTTYMaps = tty_mapper
349  battor_device_mapping.GetBattorList = battor_lister
350  find_usb_devices.GetBusNumberToDeviceTreeMap = devtree
351  battor_device_mapping.IsBattor = is_battor
352  battor_device_mapping.GetBattorSerialNumbers = battor_serials
353
354class BattorMappingTest(unittest.TestCase):
355  def tearDown(self):
356    find_usb_devices.GetAllPhysicalPortToSerialMaps = original_PPTSM
357    find_usb_devices.GetAllPhysicalPortToTTYMaps = original_PPTTM
358    battor_device_mapping.GetBattorList = original_GBL
359    find_usb_devices.GetBusNumberToDeviceTreeMap = original_GBNDM
360    battor_device_mapping.IsBattor = original_IB
361    battor_device_mapping.GetBattorSerialNumbers = original_GBSM
362
363  def test_generate_serial_map(self):
364    setup_battor_test([{1:'Phn1', 2:'Phn2', 3:'Phn3'},
365                       {1:'Bat1', 2:'Bat2', 3:'Bat3'}],
366                      [{},
367                       {1:'ttyUSB0', 2:'ttyUSB1', 3:'ttyUSB2'}],
368                      ['ttyUSB0', 'ttyUSB1', 'ttyUSB2'],
369                      ['Bat1', 'Bat2', 'Bat3'])
370    result = battor_device_mapping.GenerateSerialMap()
371    self.assertEqual(len(result), 3)
372    self.assertEqual(result['Phn1'], 'Bat1')
373    self.assertEqual(result['Phn2'], 'Bat2')
374    self.assertEqual(result['Phn3'], 'Bat3')
375
376
377if __name__ == "__main__":
378  logging.getLogger().setLevel(logging.DEBUG)
379  unittest.main(verbosity=2)
380