Generate reports from advanced test scripts.

Add functionality to parse iperf3 results, raw .pcap packet captures
(for 802.11ac VHT headers), and isostream results.

Also update wvtest to latest version from catawampus.

Change-Id: I438063e6d3a8a2f71dc62112021ade509c361cae
diff --git a/wifitables/Makefile b/wifitables/Makefile
index 9d8e033..57ffad6 100644
--- a/wifitables/Makefile
+++ b/wifitables/Makefile
@@ -16,33 +16,26 @@
 PYTHONPATH:=$(shell /bin/pwd):$(shell /bin/pwd)/../wvtest:$(PYTHONPATH)
 
 all:
-	@echo "Nothing to do."
+
+%.test: %_test.py
+	./$<
+
+runtests: \
+	$(patsubst %_test.py,%.test,$(wildcard *_test.py))
+
+lint: $(filter-out options.py,$(wildcard *.py))
+	$(GPYLINT) $^
+
+test_only: all
+	$(MAKE) runtests
 
 # Use a submake here, only because otherwise GNU make (3.81) will not print
 # an error about 'test' itself failing if one of the two sub-targets fails.
 # Without such output, 'lint' could fail long before test_only fails, and
 # the test_only output could scroll off the top of the screen, leaving the
 # misleading impression that everything tested successfully.
-test:
+test: all
 	$(MAKE) test_only lint
 
-test_only: all *_test.py
-	for d in $(filter %_test.py,$^); do \
-		echo "Testing $$d"; \
-		wvtest/wvtest.py $$d; \
-	done
-
-# For maximum parallelism, we could just have a rule that depends on %.lint
-# for all interesting files.  But gpylint takes a long time to start up, so
-# let's try to batch several files together into each instance to minimize
-# the runtime.  For added fun, gpylint has bugs if you specify files from
-# more than one directory at once, so break it out by directory.
-lint: \
-    report.lint \
-    sample.lint \
-
-%.lint: all
-	@$(GPYLINT) $*
-
 clean:
 	rm -f *~ .*~ *.pyc
diff --git a/wifitables/ifstats.py b/wifitables/ifstats.py
index 169b52c..a148f72 100644
--- a/wifitables/ifstats.py
+++ b/wifitables/ifstats.py
@@ -208,11 +208,12 @@
 def Restore(report_dir):
   """Restore an InterfaceStats cache from data on the filesystem."""
   cache = {}
-
-  for system, names in {
+  system_results = {
       'Darwin': ['ifconfig', 'airport', 'airportscan', 'systemprofiler'],
       'Linux': ['ipaddr', 'iwlink', 'iwscan'],
-      }:
+  }
+
+  for system, names in system_results.items():
     read = 0
     for name in names:
       try:
diff --git a/wifitables/ifstats_skids_test.py b/wifitables/ifstats_skids_test.py
old mode 100644
new mode 100755
index 8b58ac4..4e4e287
--- a/wifitables/ifstats_skids_test.py
+++ b/wifitables/ifstats_skids_test.py
@@ -1,15 +1,16 @@
+#!/usr/bin/python2.7
 """Tests for ifstats_skids."""
 
 import ifstats_skids
-from wvtest import *  # pylint: disable=wildcard-import
+from wvtest import wvtest
 
 
-@wvtest
+@wvtest.wvtest
 def ParseSavedStatusWireless():
-  with open('testdata/status_wireless.html') as status_wireless:
+  with open('testdata/skids/status_wireless.html') as status_wireless:
     text = status_wireless.read()
     res = ifstats_skids.ParseStatusWireless(text)
-    WVPASSEQ(res, {
+    wvtest.WVPASSEQ(res, {
         'Wireless Band': '802.11ac',
         'AP Mac Address (BSSID)': '00:26:86:F0:22:C9',
         'Bytes Transmitted': '193033536',
@@ -23,3 +24,5 @@
         'Channel': '108'
     })
 
+if __name__ == '__main__':
+  wvtest.wvtest_main()
diff --git a/wifitables/iperf.py b/wifitables/iperf.py
index 4815762..265c520 100644
--- a/wifitables/iperf.py
+++ b/wifitables/iperf.py
@@ -1,6 +1,7 @@
 """iperf: run a series of iperf tests over a wireless network."""
 
 import collections
+import glob
 import json
 import os
 import re
@@ -147,10 +148,10 @@
 def Restore(report_dir):
   """Restores an `Iperf` cache from data on the filesystem."""
   cache = {}
-  for name in ['iperf', 'iperfu']:
+  for name in glob.glob(os.path.join(report_dir, 'iperf*')):
     try:
-      with open(os.path.join(report_dir, name)) as infile:
-        cache[name] = infile.read()
+      with open(name) as infile:
+        cache[os.path.basename(name)] = infile.read()
     except IOError:
       cache[name] = ''
 
diff --git a/wifitables/isostream.py b/wifitables/isostream.py
index bb546e8..421b74c 100644
--- a/wifitables/isostream.py
+++ b/wifitables/isostream.py
@@ -1,6 +1,5 @@
-"""isostream: runs an isostream test"""
+"""isostream: runs an isostream test."""
 
-import pipes
 
 def RunIsostreamTest(run, host, bandwidth=14, time=10):
   """RunIsostreamTest runs an isostream test and returns the results."""
@@ -9,3 +8,12 @@
                    'isostream', '-b', '{:d}'.format(bandwidth), host]
   return run(isostream_cmd)
 
+
+def ParseIsostream(text):
+  # The last non-dropout line is a reasonable summary of an isostream test.
+  for line in reversed(text.splitlines()):
+    if line.find('offset=') > -1:
+      return line.strip()
+
+  return None
+
diff --git a/wifitables/report.py b/wifitables/report.py
index 67719b2..b5770c6 100755
--- a/wifitables/report.py
+++ b/wifitables/report.py
@@ -1,16 +1,19 @@
-#!/usr/bin/env python
+#!/usr/bin/python2.7
 
 """report: make a table summarizing output from one or more runs of `sample`."""
 
-from collections import Counter
+import collections
 import csv
+import datetime
 import os
 import re
 import sys
 
 import ifstats
 import iperf
+import isostream
 import options
+import wifipacket
 
 optspec = """
 report [options...] <journal>
@@ -66,7 +69,7 @@
   # assume long guard interval
   guard = 800
 
-  counter = Counter()
+  counter = collections.Counter()
   for line in outfile:
     for tok in line.split():
       if tok == '.': continue
@@ -108,14 +111,22 @@
           or (b1 <= t2 <= t1) or (b2 <= t1 <= t2))
 
 
-def ReportLine(report_dir):
+def ReportLine(report_dir, series=None):
   """Condense the output of a sample.py run into a one-line summary report."""
-  _, _, steps = os.path.basename(report_dir).split('-')
-  line = [int(steps)]
+  line = collections.OrderedDict()
+  if series:
+    line['Series'] = series
+
+  _, stamp, steps = os.path.basename(report_dir).split('-')
+  line['Time'] = datetime.datetime.fromtimestamp(float(stamp))
+  line['Steps'] = int(steps)
 
   system, cache = ifstats.Restore(report_dir)
   result = ifstats.Parse(system, cache)
-  if system == 'Darwin':
+
+  if not result.get('link'):
+    pass
+  elif system == 'Darwin':
     airport = result.get('link')
     channel, width = Channel(airport['channel'])
     shared = 0
@@ -130,69 +141,156 @@
         if Overlap(channel, width, oc, ow):
           overlap += 1
 
-    rssi = airport['agrCtlRSSI']
-    noise = airport['agrCtlNoise']
-    line += [channel, width, rssi, noise, shared, overlap - shared]
+    line.update({
+        'Channel': channel,
+        'Width': width,
+        'RSSI': airport['agrCtlRSSI'],
+        'Noise': airport['agrCtlNoise'],
+        'Shared': shared,
+        'Interfering': overlap - shared
+    })
   elif system == 'Linux':
     iwlink = result.get('link')
     signal = int(iwlink.get('signal', '0 dBm').split()[0])
-    channel = int(iwlink.get('freq'))
+    channel = int(iwlink.get('freq', '0'))
     width = 20
-    m = re.search(r'(\d+)MHz', iwlink.get('tx bitrate'), flags=re.I)
+    m = re.search(r'(\d+)MHz', iwlink.get('tx bitrate', ''), flags=re.I)
     if m:
       width = int(m.group(1))
 
     # Noise and contention not yet gathered in samples run on Linux systems.
-    line += [channel, width, signal, None, None, None]
+    line.update({
+        'Channel': channel,
+        'Width': width,
+        'RSSI': signal,
+    })
 
-  mpath = os.path.join(report_dir, 'mcs')
-  if os.path.exists(mpath):
-    with open(mpath) as mf:
-      line += ParseMCSFile(mf, width)
-  else:
-    line += [None, None]
+  try:
+    ppath = os.path.join(report_dir, 'testnetwork.pcap')
+    with open(ppath) as stream:
+      rates = [float(opt.rate) for opt, _ in wifipacket.Packetize(stream)]
+
+      # TODO(willangley): come up with a meaningful modal MCS for mixed
+      #   802.11n/802.11ac captures like we have here.
+      line['PHY'] = sum(rates)/max(len(rates), 1)
+  except IOError:
+    try:
+      mpath = os.path.join(report_dir, 'mcs')
+      with open(mpath) as mf:
+        mcs, phy = ParseMCSFile(mf, width)
+        line['MCS'] = mcs
+        line['PHY'] = phy
+    except IOError:
+      pass
 
   # If the initial ping test fails, we won't collect performance information.
   # deal with this gracefully.
   ips = iperf.Restore(report_dir)
-  for perf in [iperf.ParseIperfTCP(ips.get('iperf', '')),
-               iperf.ParseIperfUDP(ips.get('iperfu', ''))]:
-    line += [perf.get('bandwidth'), perf.get('bandwidth_unit')]
+  if 'iperf' in ips:
+    # pylint:disable=line-too-long
+    for key, perf in [('TCP BW up', iperf.ParseIperfTCP(ips.get('iperf', ''))),
+                      ('UDP BW up', iperf.ParseIperfUDP(ips.get('iperfu', '')))]:
+      line[key] = perf.get('bandwidth')
+      line['{} units'.format(key)] = perf.get('bandwidth_unit')
+  elif 'iperf3' in ips:
+    for name in (key for key in ips
+                 if key.startswith('iperf3')):
+      perf = iperf.ParseIperf3(ips[name])
+      if not perf or 'error' in perf:
+        continue
+
+      test_start = perf['start']['test_start']
+      protocol = test_start['protocol']
+      direction = 'down' if test_start['reverse'] else 'up'
+      key = '{protocol} BW {direction}'.format(protocol=protocol,
+                                               direction=direction)
+
+      if protocol == 'TCP':
+        line[key] = perf['end']['sum_received']['bits_per_second']
+      elif protocol == 'UDP':
+        line[key] = perf['end']['sum']['bits_per_second']
+      else:
+        continue
+
+      line['{} units'.format(key)] = 'bit/s'
+
+  try:
+    with open(os.path.join(report_dir, 'isostream')) as istm:
+      text = istm.read()
+      line['isostream'] = isostream.ParseIsostream(text)
+  except IOError:
+    pass
 
   return line
 
 
+def ReadJournal(jname):
+  """Read a journal, returning a series name and its data folders."""
+  jname = os.path.realpath(jname)
+  series = os.path.basename(jname)
+  if series == 'journal':
+    series = os.path.basename(os.path.dirname(jname))
+
+  folders = []
+  with open(jname) as journal:
+    for line in journal:
+      line = line.strip()
+      if line.startswith('#'):
+        continue
+      folders.append(os.path.join(os.path.dirname(jname), line))
+
+  return series, folders
+
+
+def Report(journals):
+  """Given the name of a journal file, return a list of ReportLines."""
+  report = []
+  bad = []
+
+  for jname in journals:
+    series, folders = ReadJournal(jname)
+    for folder in folders:
+      try:
+        report += [ReportLine(folder, series=series)]
+      except (TypeError, IOError) as e:
+        bad += [collections.OrderedDict(folder=folder, error=repr(e))]
+
+  return report, bad
+
+
+def WriteReport(lines):
+  """Write a network testing report in .tsv format to stdout."""
+  # include every field we can write in the header row
+  header = ['Series', 'Time', 'Steps', 'Channel', 'Width', 'RSSI', 'Noise',
+            'Shared', 'Interfering', 'MCS', 'PHY', 'TCP BW up',
+            'TCP BW up units', 'UDP BW up', 'UDP BW up units', 'TCP BW down',
+            'TCP BW down units', 'UDP BW down', 'UDP BW down units',
+            'isostream']
+
+  writer = csv.DictWriter(sys.stdout, header, dialect=csv.excel_tab)
+  writer.writeheader()
+  writer.writerows(lines)
+
+
 def main():
   o = options.Options(optspec)
   (opt, _, extra) = o.parse(sys.argv[1:])
-  if len(extra) > 1:
-    o.fatal('expected at most one journal name.')
 
-  LoadNRates()
-  LoadChannels()
-
-  lines = []
   if opt.report_dir:
-    lines += [ReportLine(opt.report_dir)]
+    report = [ReportLine(opt.report_dir)]
+  elif extra:
+    report, bad = Report(extra)
+    if bad:
+      writer = csv.DictWriter(sys.stdout, bad[0].keys(), dialect=csv.excel_tab)
+      writer.writeheader()
+      writer.writerows(bad)
+      print
 
-  if extra:
-    for jname in extra[:1]:
-      jname = os.path.realpath(jname)
-      with open(jname) as journal:
-        for line in journal:
-          lines += [ReportLine(os.path.join(os.path.dirname(jname),
-                                            line.strip()))]
-
-  if len(lines) < 1:
-    o.fatal("didn't find any samples. did you supply at least one report dir"
+  if len(report) < 1:
+    o.fatal("Didn't find any samples. Did you supply at least one report dir"
             ' or journal?')
 
-  header = ['Steps', 'Channel', 'Width', 'RSSI', 'Noise', 'Shared',
-            'Interfering', 'MCS', 'PHY', 'TCP BW', '(Units)', 'UDP BW',
-            '(Units)']
-  writer = csv.writer(sys.stdout, delimiter='\t', quoting=csv.QUOTE_MINIMAL)
-  writer.writerow(header)
-  writer.writerows(lines)
+  WriteReport(report)
 
 
 if __name__ == '__main__':
diff --git a/wifitables/report_test.py b/wifitables/report_test.py
old mode 100644
new mode 100755
index 6ccb832..a7d5a87
--- a/wifitables/report_test.py
+++ b/wifitables/report_test.py
@@ -1,116 +1,147 @@
+#!/usr/bin/python2.7 -S
 """Tests for report."""
 
-import os
-
 import ifstats
 import report
-from wvtest import *  # pylint: disable=wildcard-import
+from wvtest import wvtest
 
 
-@wvtest
+@wvtest.wvtest
 def LoadExternalData():
   print
   report.LoadNRates()
-  WVPASS(report.nrates)
+  wvtest.WVPASS(report.nrates)
 
   idx = (0, 20, 800)
   print 'Testing MCS rate in file', idx
-  WVPASSEQ(report.nrates[idx], 6.5)
+  wvtest.WVPASSEQ(report.nrates[idx], 6.5)
 
   idx = (25, 40, 400)
   print 'Testing computed MCS rate', idx
-  WVPASSEQ(report.nrates[idx], 120)
+  wvtest.WVPASSEQ(report.nrates[idx], 120)
 
   print
   report.LoadChannels()
-  WVPASS(report.channels)
+  wvtest.WVPASS(report.channels)
 
   ch = 5
   print 'Testing 2.4GHz channel', ch
-  WVPASSEQ(report.channels[ch], 2432)
+  wvtest.WVPASSEQ(report.channels[ch], 2432)
 
   ch = 149
   print 'Testing 5GHz channel', ch
-  WVPASSEQ(report.channels[ch], 5745)
+  wvtest.WVPASSEQ(report.channels[ch], 5745)
 
 
-@wvtest
+@wvtest.wvtest
 def TVBoxReport():
   rpt = 'testdata/wifi-1424739295.41-0010'
-  print 'Checking IP address'
+
+  print 'Restoring report'
   system, cache = ifstats.Restore(rpt)
+  wvtest.WVPASSEQ(system, 'Linux')
+
+  print 'Checking IP address'
   addrmap = ifstats.ParseIpAddr(cache['ipaddr'])
-  WVPASSEQ(addrmap.get('lo'), '127.0.0.1/32')
-  WVPASSEQ(addrmap.get('wcli0'), '192.168.1.222/24')
+  wvtest.WVPASSEQ(addrmap.get('lo'), '127.0.0.1/32')
+  wvtest.WVPASSEQ(addrmap.get('wcli0'), '192.168.1.222/24')
 
   print 'Checking for link information'
   data = ifstats.ParseIwLink(cache['iwlink'])
-  WVPASSEQ(data.get('SSID'), 'GSAFNS1441P0208_TestWifi')
-  WVPASSEQ(data.get('BSSID'), 'f4:f5:e8:80:f3:d0')
+  wvtest.WVPASSEQ(data.get('SSID'), 'GSAFNS1441P0208_TestWifi')
+  wvtest.WVPASSEQ(data.get('BSSID'), 'f4:f5:e8:80:f3:d0')
 
-  line = report.ReportLine(rpt)
-  print ('Checking report. Implemented measures: steps, channel, width, rssi, '
-         'TCP performance, UDP performance')
-  print line
+  got = report.ReportLine(rpt)
+  want = {
+      'Steps': 10,
+      'Channel': 5745,
+      'Width': 80,
+      'RSSI': -39,
+      'TCP BW up': 92.1,
+      'TCP BW up units': 'Mbits/sec',
+      'UDP BW up': 91.2,
+      'UDP BW up units': 'Mbits/sec',
+  }
 
-  (steps, ch, width, rssi, _, _, _, _, _,
-   tcperf, tcpbw, udperf, udpbw) = line
+  print 'Checking report.'
+  for key, value in want.items():
+    print key
+    wvtest.WVPASSEQ(got[key], value)
+    del got[key]
 
-  WVPASSEQ(steps, 10)
-  WVPASSEQ(ch, 5745)
-  WVPASSEQ(width, 80)
-  WVPASSEQ(rssi, -39)
-  WVPASSEQ(tcperf, 92.1)
-  WVPASSEQ(tcpbw, 'Mbits/sec')
-  WVPASSEQ(udperf, 91.2)
-  WVPASSEQ(udpbw, 'Mbits/sec')
+  print
+  print 'Not checked:', got
 
 
-@wvtest
+@wvtest.wvtest
 def MacbookReport():
   rpt = 'testdata/wifi-1424744066.47-0010'
+  got = report.ReportLine(rpt)
+  want = {
+      'Steps': 10,
+      'Channel': 5755,
+      'Width': 40,
+      'RSSI': -29,
+      'Noise': -90,
+      'Shared': 17,
+      'Interfering': 10,
+      'MCS': 21,
+      'TCP BW up': 196.0,
+      'TCP BW up units': 'Mbits/sec',
+      'UDP BW up': 260.0,
+      'UDP BW up units': 'Mbits/sec',
+  }
+
+  print 'Checking report.'
+  for key, value in want.items():
+    print key
+    wvtest.WVPASSEQ(got[key], value)
+    del got[key]
+
+  print 'PHY rate'
+  wvtest.WVPASS(abs(got['PHY'] - 340.9) < 0.1)
+  del got['PHY']
+
+  print
+  print 'Not checked:', got
+
+
+@wvtest.wvtest
+def SkidReport():
+  rpt = 'testdata/wifi-1431102603.48-0103'
+
+  print 'Checking subset of skid report.'
   line = report.ReportLine(rpt)
-
-  print ('Checking report. Implemented measures: steps, channel, width, rssi, '
-         'noise, devices on channel, off channel, MCS, PHY rate, '
-         'TCP performance, UDP performance.')
-  print line
-
-  (steps, ch, width, rssi, noise, shared, conflict,
-   mcs, phy, tcperf, tcpbw, udperf, udpbw) = line
-
-  WVPASSEQ(steps, 10)
-  WVPASSEQ(ch, 5755)
-  WVPASSEQ(width, 40)
-  WVPASSEQ(rssi, -29)
-  WVPASSEQ(noise, -90)
-  WVPASSEQ(shared, 17)
-  WVPASSEQ(conflict, 10)
-  WVPASSEQ(mcs, 21)
-  WVPASS(abs(phy - 340.9) < 0.1)
-  WVPASSEQ(tcperf, 196.0)
-  WVPASSEQ(tcpbw, 'Mbits/sec')
-  WVPASSEQ(udperf, 260.0)
-  WVPASSEQ(udpbw, 'Mbits/sec')
+  wvtest.WVPASSEQ(line['isostream'],
+                  '50.638s 14Mbps offset=0.003s disconn=0/0.000s '
+                  'drops=2/0.009s/-0.049s')
 
 
-@wvtest
+@wvtest.wvtest
 def ReportWithHyphensInPath():
   rpt = 'testdata/this-failed-before/wifi-1425669615.33-0010'
   try:
     _ = report.ReportLine(rpt)
   except ValueError:
-    WVFAIL('Failed to read report directory with hyphens in path.')
+    wvtest.WVFAIL('Failed to read report directory with hyphens in path.')
     return
 
-  WVPASS('Report with hyphens in path read successfully.')
+  wvtest.WVPASS('Report with hyphens in path read successfully.')
 
 
-@wvtest
+@wvtest.wvtest
 def ReportWithoutIperfFiles():
   try:
     rl = report.ReportLine('testdata/wifi-1426545351.85-0085')
-    WVPASSEQ(rl[-4:], [None, None, None, None])
+    for key, value in rl.items():
+      if key.startswith('TCP BW') or key.startswith('UDP BW'):
+        wvtest.WVFAIL("Report included unexpected field: report['{}'] = {}"
+                      .format(key, value))
   except IOError:
-    WVFAIL('Failed to read report without iperf output files.')
+    wvtest.WVFAIL('Failed to read report without iperf output files.')
 
+  wvtest.WVPASS('Report without iperf files succeeded.')
+
+
+if __name__ == '__main__':
+  wvtest.wvtest_main()
diff --git a/wvtest/wvtest.py b/wvtest/wvtest.py
index d424c55..8f2564c 100755
--- a/wvtest/wvtest.py
+++ b/wvtest/wvtest.py
@@ -6,23 +6,41 @@
 #       See the included file named LICENSE for license information.
 #       You can get wvtest from: http://github.com/apenwarr/wvtest
 #
-import traceback
+import atexit
+import inspect
 import os
 import re
 import sys
+import traceback
 
-if __name__ != "__main__":   # we're imported as a module
+# NOTE
+# Why do we do we need the "!= main" check?  Because if you run
+# wvtest.py as a main program and it imports your test files, then
+# those test files will try to import the wvtest module recursively.
+# That actually *works* fine, because we don't run this main program
+# when we're imported as a module.  But you end up with two separate
+# wvtest modules, the one that gets imported, and the one that's the
+# main program.  Each of them would have duplicated global variables
+# (most importantly, wvtest._registered), and so screwy things could
+# happen.  Thus, we make the main program module *totally* different
+# from the imported module.  Then we import wvtest (the module) into
+# wvtest (the main program) here and make sure to refer to the right
+# versions of global variables.
+#
+# All this is done just so that wvtest.py can be a single file that's
+# easy to import into your own applications.
+if __name__ != '__main__':   # we're imported as a module
     _registered = []
     _tests = 0
     _fails = 0
 
-    def wvtest(func):
-        """ Use this decorator (@wvtest) in front of any function you want to run
-            as part of the unit test suite.  Then run:
-                python wvtest.py path/to/yourtest.py
-            to run all the @wvtest functions in that file.
+    def wvtest(func, innerfunc=None):
+        """ Use this decorator (@wvtest) in front of any function you want to
+            run as part of the unit test suite.  Then run:
+                python wvtest.py path/to/yourtest.py [other test.py files...]
+            to run all the @wvtest functions in the given file(s).
         """
-        _registered.append(func)
+        _registered.append((func, innerfunc or func))
         return func
 
 
@@ -40,8 +58,8 @@
         sys.stdout.flush()
 
 
-    def _check(cond, msg = 'unknown', tb = None):
-        if tb == None: tb = traceback.extract_stack()[-3]
+    def _check(cond, msg, xdepth):
+        tb = traceback.extract_stack()[-3 - xdepth]
         if cond:
             _result(msg, tb, 'ok')
         else:
@@ -49,113 +67,186 @@
         return cond
 
 
-    def _code():
-        (filename, line, func, text) = traceback.extract_stack()[-3]
-        text = re.sub(r'^\w+\((.*)\)(\s*#.*)?$', r'\1', text);
+    def _code(xdepth):
+        (filename, line, func, text) = traceback.extract_stack()[-3 - xdepth]
+        text = re.sub(r'^[\w\.]+\((.*)\)(\s*#.*)?$', r'\1', str(text));
         return text
 
 
-    def WVPASS(cond = True):
+    def WVPASS(cond = True, xdepth = 0):
         ''' Counts a test failure unless cond is true. '''
-        return _check(cond, _code())
+        return _check(cond, _code(xdepth), xdepth)
 
-    def WVFAIL(cond = True):
+    def WVFAIL(cond = True, xdepth = 0):
         ''' Counts a test failure  unless cond is false. '''
-        return _check(not cond, 'NOT(%s)' % _code())
+        return _check(not cond, 'NOT(%s)' % _code(xdepth), xdepth)
 
-    def WVPASSEQ(a, b):
+    def WVPASSIS(a, b, xdepth = 0):
+        ''' Counts a test failure unless a is b. '''
+        return _check(a is b, '%s is %s' % (repr(a), repr(b)), xdepth)
+
+    def WVPASSISNOT(a, b, xdepth = 0):
+        ''' Counts a test failure unless a is not b. '''
+        return _check(a is not b, '%s is not %s' % (repr(a), repr(b)), xdepth)
+
+    def WVPASSEQ(a, b, xdepth = 0):
         ''' Counts a test failure unless a == b. '''
-        return _check(a == b, '%s == %s' % (repr(a), repr(b)))
+        return _check(a == b, '%s == %s' % (repr(a), repr(b)), xdepth)
 
-    def WVPASSNE(a, b):
+    def WVPASSNE(a, b, xdepth = 0):
         ''' Counts a test failure unless a != b. '''
-        return _check(a != b, '%s != %s' % (repr(a), repr(b)))
+        return _check(a != b, '%s != %s' % (repr(a), repr(b)), xdepth)
 
-    def WVPASSLT(a, b):
+    def WVPASSLT(a, b, xdepth = 0):
         ''' Counts a test failure unless a < b. '''
-        return _check(a < b, '%s < %s' % (repr(a), repr(b)))
+        return _check(a < b, '%s < %s' % (repr(a), repr(b)), xdepth)
 
-    def WVPASSLE(a, b):
+    def WVPASSLE(a, b, xdepth = 0):
         ''' Counts a test failure unless a <= b. '''
-        return _check(a <= b, '%s <= %s' % (repr(a), repr(b)))
+        return _check(a <= b, '%s <= %s' % (repr(a), repr(b)), xdepth)
 
-    def WVPASSGT(a, b):
+    def WVPASSGT(a, b, xdepth = 0):
         ''' Counts a test failure unless a > b. '''
-        return _check(a > b, '%s > %s' % (repr(a), repr(b)))
+        return _check(a > b, '%s > %s' % (repr(a), repr(b)), xdepth)
 
-    def WVPASSGE(a, b):
+    def WVPASSGE(a, b, xdepth = 0):
         ''' Counts a test failure unless a >= b. '''
-        return _check(a >= b, '%s >= %s' % (repr(a), repr(b)))
+        return _check(a >= b, '%s >= %s' % (repr(a), repr(b)), xdepth)
 
-    def WVEXCEPT(etype, func, *args, **kwargs):
+    def WVPASSNEAR(a, b, places = 7, delta = None, xdepth = 0):
+        ''' Counts a test failure unless a ~= b. '''
+        if delta:
+            return _check(abs(a - b) <= abs(delta),
+                          '%s ~= %s' % (repr(a), repr(b)), xdepth)
+        else:
+            return _check(round(a, places) == round(b, places),
+                          '%s ~= %s' % (repr(a), repr(b)), xdepth)
+
+    def WVPASSFAR(a, b, places = 7, delta = None, xdepth = 0):
+        ''' Counts a test failure unless a ~!= b. '''
+        if delta:
+            return _check(abs(a - b) > abs(delta),
+                          '%s ~= %s' % (repr(a), repr(b)), xdepth)
+        else:
+            return _check(round(a, places) != round(b, places),
+                          '%s ~= %s' % (repr(a), repr(b)), xdepth)
+
+    def _except_report(cond, code, xdepth):
+        return _check(cond, 'EXCEPT(%s)' % code, xdepth + 1)
+
+    class _ExceptWrapper(object):
+        def __init__(self, etype, xdepth):
+            self.etype = etype
+            self.xdepth = xdepth
+            self.code = None
+
+        def __enter__(self):
+          self.code = _code(self.xdepth)
+
+        def __exit__(self, etype, value, traceback):
+            if etype == self.etype:
+                _except_report(True, self.code, self.xdepth)
+                return 1  # success, got the expected exception
+            elif etype is None:
+                _except_report(False, self.code, self.xdepth)
+                return 0
+            else:
+                _except_report(False, self.code, self.xdepth)
+
+    def _WVEXCEPT(etype, xdepth, func, *args, **kwargs):
+        if func:
+            code = _code(xdepth + 1)
+            try:
+                func(*args, **kwargs)
+            except etype, e:
+                return _except_report(True, code, xdepth + 1)
+            except:
+                _except_report(False, code, xdepth + 1)
+                raise
+            else:
+                return _except_report(False, code, xdepth + 1)
+        else:
+            return _ExceptWrapper(etype, xdepth)
+
+    def WVEXCEPT(etype, func=None, *args, **kwargs):
         ''' Counts a test failure unless func throws an 'etype' exception.
             You have to spell out the function name and arguments, rather than
             calling the function yourself, so that WVEXCEPT can run before
             your test code throws an exception.
         '''
-        try:
-            func(*args, **kwargs)
-        except etype, e:
-            return _check(True, 'EXCEPT(%s)' % _code())
-        except:
-            _check(False, 'EXCEPT(%s)' % _code())
-            raise
-        else:
-            return _check(False, 'EXCEPT(%s)' % _code())
+        return _WVEXCEPT(etype, 0, func, *args, **kwargs)
 
-else:  # we're the main program
-    # NOTE
-    # Why do we do this in such a convoluted way?  Because if you run
-    # wvtest.py as a main program and it imports your test files, then
-    # those test files will try to import the wvtest module recursively.
-    # That actually *works* fine, because we don't run this main program
-    # when we're imported as a module.  But you end up with two separate
-    # wvtest modules, the one that gets imported, and the one that's the
-    # main program.  Each of them would have duplicated global variables
-    # (most importantly, wvtest._registered), and so screwy things could
-    # happen.  Thus, we make the main program module *totally* different
-    # from the imported module.  Then we import wvtest (the module) into
-    # wvtest (the main program) here and make sure to refer to the right
-    # versions of global variables.
-    #
-    # All this is done just so that wvtest.py can be a single file that's
-    # easy to import into your own applications.
-    import wvtest
 
-    def _runtest(modname, fname, f):
+    def _check_unfinished():
+        if _registered:
+            for func, innerfunc in _registered:
+                print 'WARNING: not run: %r' % (innerfunc,)
+            WVFAIL('wvtest_main() not called')
+        if _fails:
+            sys.exit(1)
+
+    atexit.register(_check_unfinished)
+
+
+def _run_in_chdir(path, func, *args, **kwargs):
+    oldwd = os.getcwd()
+    oldpath = sys.path
+    try:
+        if path: os.chdir(path)
+        sys.path += [path, os.path.split(path)[0]]
+        return func(*args, **kwargs)
+    finally:
+        os.chdir(oldwd)
+        sys.path = oldpath
+
+
+def _runtest(fname, f, innerfunc):
+    import wvtest as _wvtestmod
+    mod = inspect.getmodule(innerfunc)
+    relpath = os.path.relpath(mod.__file__, os.getcwd()).replace('.pyc', '.py')
+    print
+    print 'Testing "%s" in %s:' % (fname, relpath)
+    sys.stdout.flush()
+    try:
+        _run_in_chdir(os.path.split(mod.__file__)[0], f)
+    except Exception, e:
         print
-        print 'Testing "%s" in %s.py:' % (fname, modname)
-        sys.stdout.flush()
-        try:
-            f()
-        except Exception, e:
-            print
-            print traceback.format_exc()
-            tb = sys.exc_info()[2]
-            wvtest._result(e, traceback.extract_tb(tb)[1], 'EXCEPTION')
+        print traceback.format_exc()
+        tb = sys.exc_info()[2]
+        _wvtestmod._result(repr(e), traceback.extract_tb(tb)[-1], 'EXCEPTION')
 
-    # main code
-    for modname in sys.argv[1:]:
+
+def _run_registered_tests():
+    import wvtest as _wvtestmod
+    while _wvtestmod._registered:
+        func, innerfunc = _wvtestmod._registered.pop(0)
+        _runtest(innerfunc.func_name, func, innerfunc)
+        print
+
+
+def wvtest_main(extra_testfiles=[]):
+    import wvtest as _wvtestmod
+    _run_registered_tests()
+    for modname in extra_testfiles:
         if not os.path.exists(modname):
             print 'Skipping: %s' % modname
             continue
         if modname.endswith('.py'):
             modname = modname[:-3]
         print 'Importing: %s' % modname
-        wvtest._registered = []
-        oldwd = os.getcwd()
-        oldpath = sys.path
-        try:
-            path, mod = os.path.split(os.path.abspath(modname))
-            os.chdir(path)
-            sys.path += [path, os.path.split(path)[0]]
-            mod = __import__(modname.replace(os.path.sep, '.'), None, None, [])
-            for t in wvtest._registered:
-                _runtest(modname, t.func_name, t)
-                print
-        finally:
-            os.chdir(oldwd)
-            sys.path = oldpath
-
+        path, mod = os.path.split(os.path.abspath(modname))
+        nicename = modname.replace(os.path.sep, '.')
+        while nicename.startswith('.'):
+            nicename = modname[1:]
+        _run_in_chdir(path, __import__, nicename, None, None, [])
+        _run_registered_tests()
     print
-    print 'WvTest: %d tests, %d failures.' % (wvtest._tests, wvtest._fails)
+    print 'WvTest: %d tests, %d failures.' % (_wvtestmod._tests,
+                                              _wvtestmod._fails)
+
+
+if __name__ == '__main__':
+    import wvtest as _wvtestmod
+    sys.modules['wvtest'] = _wvtestmod
+    sys.modules['wvtest.wvtest'] = _wvtestmod
+    wvtest_main(sys.argv[1:])