[project @ robey@lag.net-20060731070137-571efc087c20347a]

clean up test.py a bit and allow filtering from the command line

test.py

@@ -22,7 +22,10 @@
 do the unit tests!
 """
 
-import sys, os, unittest
+import os
+import re
+import sys
+import unittest
 from optparse import OptionParser
 import paramiko
 
@@ -46,6 +49,30 @@ default_user = os.environ.get('USER', 'nobody')
 default_keyfile = os.path.join(os.environ.get('HOME', '/'), '.ssh/id_rsa')
 default_passwd = None
 
+
+def iter_suite_tests(suite):
+    """Return all tests in a suite, recursing through nested suites"""
+    for item in suite._tests:
+        if isinstance(item, unittest.TestCase):
+            yield item
+        elif isinstance(item, unittest.TestSuite):
+            for r in iter_suite_tests(item):
+                yield r
+        else:
+            raise Exception('unknown object %r inside test suite %r'
+                            % (item, suite))
+
+
+def filter_suite_by_re(suite, pattern):
+    result = unittest.TestSuite()
+    filter_re = re.compile(pattern)
+    for test in iter_suite_tests(suite):
+        if filter_re.search(test.id()):
+            result.addTest(test)
+    return result
+
+
+def main():
     parser = OptionParser('usage: %prog [options]')
     parser.add_option('--verbose', action='store_true', dest='verbose', default=False,
                       help='verbose display (one line per test)')
@@ -75,8 +102,6 @@ parser.add_option('-P', '--sftp-passwd', dest='password', type='string', default
                       help='[with -R] (optional) password to unlock the private key for remote sftp tests')
     
     options, args = parser.parse_args()
-if len(args) > 0:
-    parser.error('unknown argument(s)')
     
     # setup logging
     paramiko.util.log_to_file('test.log')
@@ -109,4 +134,13 @@ if options.use_big_file:
     verbosity = 1
     if options.verbose:
         verbosity = 2
-unittest.TextTestRunner(verbosity=verbosity).run(suite)
+
+    runner = unittest.TextTestRunner(verbosity=verbosity)
+    if len(args) > 0:
+        filter = '|'.join(args)
+        suite = filter_suite_by_re(suite, filter)
+    runner.run(suite)
+
+
+if __name__ == '__main__':
+    main()