wifitables/iperf.py

"""iperf: run a series of iperf tests over a wireless network."""

import os
import pipes
import re
import sys

DEVNULL = open(os.devnull, 'wb')


class Iperf(object):
  """Iperf collects `iperf` measurements across a wireless network."""

  def __init__(self, runner=None, bind=None):
    self.run = runner
    self.bind = bind
    self.cache = {}

  def Ping(self, host):
    line = 'ping -c 1 {}'.format(pipes.quote(host))
    return self.run(line).return_code

  def _Iperf(self, host, udp=False, bandwidth=20):
    """Run iperf against host and return string containing stdout from run."""
    line = 'iperf -c {}'
    args = [host]
    if udp:
      line += ' -u -b {}'
      args += [str(bandwidth * 1000000)]
    if self.bind:
      line += ' -B {}'
      args += [self.bind]

    return self.run(line.format(*[pipes.quote(arg) for arg in args]))

  def IperfTCP(self, host='127.0.0.1'):
    return self._Iperf(host)

  def IperfUDP(self, host='127.0.0.1', bandwidth=20):
    return self._Iperf(host, udp=True, bandwidth=bandwidth)

  def RunTestSeries(self, host):
    """RunTestSeries runs iperf tests and returns their results.

    Args:
      host: string containing the hostname to run tests against
    Returns:
      a list of files; each file contains output for one test
    """
    outputs = {}
    status = self.Ping(host)
    if not status:
      it = self.IperfTCP(host)
      outputs['iperf'] = it

      # Empirically about 1.25x more packets make it through in UDP than TCP.
      # Try to saturate the channel by sending a bit more than that over UDP.
      bandwidth = ParseIperfTCP(it).get('bandwidth', 0.01)
      outputs['iperfu'] = self.IperfUDP(host, bandwidth=bandwidth * 1.5)
    else:
      print >> sys.stderr, ('Could not ping destination host {0}; '
                            'skipping performance tests').format(host)

    return outputs

def _ParseIperf(text, udp=False):
  """Parse summary line written by an `iperf` run into a Python dict."""
  pattern = (r'\[(.{3})\]\s+(?P<interval>.*?sec)\s+'
             r'(?P<transfer>.*?Bytes|bits)'
             r'\s+(?P<bandwidth>.*?/sec)')
  if udp:
    pattern += r'\s+(?P<jitter>.*?s)\s+(?P<datagrams>.*)'

  iperf_re = re.compile(pattern)

  for line in text.splitlines():
    match = iperf_re.match(line)
    if match:
      iperf = match.groupdict()
      bval, bunit = iperf['bandwidth'].split()
      iperf['bandwidth'] = float(bval)
      iperf['bandwidth_unit'] = bunit

      tval, tunit = iperf['transfer'].split()
      iperf['transfer'] = float(tval)
      iperf['transfer_unit'] = tunit
      return iperf

  return {}


def ParseIperfTCP(text):
  # sample line: [  4]  0.0-10.0 sec   245 MBytes   206 Mbits/sec
  return _ParseIperf(text)


def ParseIperfUDP(text):
  # pylint: disable=line-too-long
  # sample line: [  5]  0.0-10.0 sec  1.25 MBytes  1.05 Mbits/sec   0.593 ms    0/  893 (0%)
  return _ParseIperf(text, udp=True)


def Restore(report_dir):
  """Restores an `Iperf` cache from data on the filesystem."""
  cache = {}
  for name in ['iperf', 'iperfu']:
    ipath = os.path.join(report_dir, name)
    if os.path.exists(ipath):
      with open(ipath) as infile:
        cache[name] = infile.read()
    else:
      cache[name] = ''

  return cache