Repository: google/OpenSK
File: tools/vendor_hid_test.py (Python)

Code before:
    sleep()
    return r

  def ping_continue(self, num, byte=0x88) -> int:
    continue_packet = [0] + list(self.cid) + [num] + [byte] * 59
    if len(continue_packet) != _SEND_DATA_SIZE:
      raise Exception(f'Expected packet to be {_SEND_DATA_SIZE} but was '
                      '{len(continue_packet)}')
    r = self.dev.write(bytes(continue_packet))
    sleep()
    return r

  def read_and_print(self) -> int:
    d = self.dev.read(_PACKET_SIZE, 2000)
    self.rx_packets.append(d)
    sleep()
    return len(d)

  def get_received_data(self):
    """This combines the data from the received packets, to match the ping
packets sent."""
    d = b''
    if len(self.rx_packets) < _PACKETS:
      raise Exception(f'Insufficent packets received - want {_PACKETS}'
                      ', got {len(self.rx_packets)}')
    d += self.rx_packets[-_PACKETS][7:]
    for p in self.rx_packets[-_PACKETS + 1:]:
      d += p[5:]
    return d


def get_devices(usage_page):
  for device in hid.enumerate(_OPENSK_VID, _OPENSK_PID):
    if device['usage_page'] == usage_page:
      yield device


class VendorHid(unittest.TestCase):
  """Tests for the Vendor HID interface."""

  @classmethod
  def setUpClass(cls):
    cls.fido_hid = cls.get_device(_FIDO_USAGE_PAGE)
    cls.vendor_hid = cls.get_device(_VENDOR_USAGE_PAGE)

  @classmethod
  def get_device(cls, usage_page):
    devices = list(get_devices(usage_page))
    if len(devices) != 1:
      raise Exception(f'Found {len(devices)} devices')
    return HidDevice(devices[0])

Review comment (suggestion):
`-_PACKETS` could simply be index 0 if we consumed the returned packages. They are not accessed afterwards, so should be safe to delete after returning the message. Also means we could check if we have any unexpected extra packets. And move the number of expected packets to the outside, as indicated above.

Code after:
    assert len(ping_packet) == _SEND_DATA_SIZE, (
        f'Expected packet to be {_SEND_DATA_SIZE} but was {len(ping_packet)}')

    r = self.dev.write(bytes(ping_packet))
    sleep()
    return r

  def ping_continue(self, num, byte=0x88) -> int:
    continue_packet = [0] + list(self.cid) + [num] + [byte] * 59
    assert len(continue_packet) == _SEND_DATA_SIZE, (
        f'Expected packet to be {_SEND_DATA_SIZE} '
        'but was {len(continue_packet)}')
    r = self.dev.write(bytes(continue_packet))
    sleep()
    return r

  def read_and_print(self) -> int:
    d = self.dev.read(_PACKET_SIZE, 2000)
    self.rx_packets.append(d)
    sleep()
    return len(d)

  def get_received_data(self) -> bytes:
    """This combines the data from the received packets, to match the ping
packets sent."""
    d = b''
    d += self.rx_packets.pop(0)[7:]
    for p in self.rx_packets:
      d += p[5:]
    # And clear the packets received to ensure consistency.
    self.rx_packets = []
    return d


def get_devices(usage_page) -> Iterable[Dict]:
  for device in hid.enumerate(_OPENSK_VID, _OPENSK_PID):
    if device['usage_page'] == usage_page:
      yield device


class HidInterfaces(unittest.TestCase):
  """Tests for the Vendor and FIDO HID interfaces."""

  @classmethod
  def setUpClass(cls):
    cls.fido_hid = cls.get_device(_FIDO_USAGE_PAGE)
    cls.vendor_hid = cls.get_device(_VENDOR_USAGE_PAGE)

  @classmethod
  def get_device(cls, usage_page) -> HidDevice:
    devices = list(get_devices(usage_page))