Passing arguments (for argparse) with unittest discover

Question

foo is a Python project with deep directory nesting, including ~30 unittest files in various subdirectories. Within foo's setup.py, I've added a custom "test" command internally running

 python -m unittest discover foo '*test.py'

Note that this uses unittest's discovery mode.

---

Since some of the tests are extremely slow, I've recently decided that tests should have "levels". The answer to this question explained very well how to get unittest and argparse to work well with each other. So now, I can run an individual unittest file, say foo/bar/_bar_test.py, with

python foo/bar/_bar_test.py --level=3

and only level-3 tests are run.

The problem is that I can't figure out how to pass the custom flag (in this case "--level=3" using discover. Everything I try fails, e.g.:

$ python -m unittest discover --level=3 foo '*test.py'
Usage: python -m unittest discover [options]

python -m unittest discover: error: no such option: --level

$ python -m --level=3 unittest discover foo '*test.py'
/usr/bin/python: No module named --level=3

How can I pass --level=3 to the individual unittests? If possible, I'd like to avoid dividing different-level tests to different files.

Bounty Edit

The pre-bounty (fine) solution suggests using system environment variables. This is not bad, but I'm looking for something cleaner.

Changing the multiple-file test runner (i.e., python -m unittest discover foo '*test.py') to something else is fine, as long as:

1. It allows generating a single report for multiple-file unittests.
2. It can somehow support multiple test levels (either using the technique in the question, or using some other different mechanism).

Answer 1

This doesn't pass args using unittest discover, but it accomplishes what you are trying to do.

This is leveltest.py. Put it somewhere in the module search path (maybe current directory or site-packages):

import argparse
import sys
import unittest

# this part copied from unittest.__main__.py
if sys.argv[0].endswith("__main__.py"):
    import os.path
    # We change sys.argv[0] to make help message more useful
    # use executable without path, unquoted
    # (it's just a hint anyway)
    # (if you have spaces in your executable you get what you deserve!)
    executable = os.path.basename(sys.executable)
    sys.argv[0] = executable + " -m leveltest"
    del os

def _id(obj):
    return obj

# decorator that assigns test levels to test cases (classes and methods)
def level(testlevel):
    if unittest.level < testlevel:
        return unittest.skip("test level too low.")
    return _id

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--level', type=int, default=3)
    ns, args = parser.parse_known_args(namespace=unittest)
    return ns, sys.argv[:1] + args

if __name__ == "__main__":
    ns, remaining_args = parse_args()

    # this invokes unittest when leveltest invoked with -m flag like:
    #    python -m leveltest --level=2 discover --verbose
    unittest.main(module=None, argv=remaining_args)

Here is how you use it in an example testproject.py file:

import unittest
import leveltest

# This is needed before any uses of the @leveltest.level() decorator
#   to parse the "--level" command argument and set the test level when 
#   this test file is run directly with -m
if __name__ == "__main__":
    ns, remaining_args = leveltest.parse_args()

@leveltest.level(2)
class TestStringMethods(unittest.TestCase):

    @leveltest.level(5)
    def test_upper(self):
        self.assertEqual('foo'.upper(), 'FOO')

    @leveltest.level(3)
    def test_isupper(self):
        self.assertTrue('FOO'.isupper())
        self.assertFalse('Foo'.isupper())

    @leveltest.level(4)
    def test_split(self):
        s = 'hello world'
        self.assertEqual(s.split(), ['hello', 'world'])
        # check that s.split fails when the separator is not a string
        with self.assertRaises(TypeError):
            s.split(2)

if __name__ == '__main__':
    # this invokes unittest when this file is executed with -m
    unittest.main(argv=remaining_args)

You can then run tests by running testproject.py directly, like:

~roottwo\projects> python testproject.py --level 2 -v
test_isupper (__main__.TestStringMethods) ... skipped 'test level too low.'
test_split (__main__.TestStringMethods) ... skipped 'test level too low.'
test_upper (__main__.TestStringMethods) ... skipped 'test level too low.'

----------------------------------------------------------------------
Ran 3 tests in 0.000s

OK (skipped=3)

~roottwo\projects> python testproject.py --level 3 -v
test_isupper (__main__.TestStringMethods) ... ok
test_split (__main__.TestStringMethods) ... skipped 'test level too low.'
test_upper (__main__.TestStringMethods) ... skipped 'test level too low.'

----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK (skipped=2)

~roottwo\projects> python testproject.py --level 4 -v
test_isupper (__main__.TestStringMethods) ... ok
test_split (__main__.TestStringMethods) ... ok
test_upper (__main__.TestStringMethods) ... skipped 'test level too low.'

----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK (skipped=1)

~roottwo\projects> python testproject.py --level 5 -v
test_isupper (__main__.TestStringMethods) ... ok
test_split (__main__.TestStringMethods) ... ok
test_upper (__main__.TestStringMethods) ... ok

----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK

By using unittest discovery like this:

~roottwo\projects> python -m leveltest --level 2 -v
test_isupper (testproject.TestStringMethods) ... skipped 'test level too low.'
test_split (testproject.TestStringMethods) ... skipped 'test level too low.'
test_upper (testproject.TestStringMethods) ... skipped 'test level too low.'

----------------------------------------------------------------------
Ran 3 tests in 0.003s

OK (skipped=3)

~roottwo\projects> python -m leveltest --level 3 discover -v
test_isupper (testproject.TestStringMethods) ... ok
test_split (testproject.TestStringMethods) ... skipped 'test level too low.'
test_upper (testproject.TestStringMethods) ... skipped 'test level too low.'

----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK (skipped=2)

~roottwo\projects> python -m leveltest --level 4 -v
test_isupper (testproject.TestStringMethods) ... ok
test_split (testproject.TestStringMethods) ... ok
test_upper (testproject.TestStringMethods) ... skipped 'test level too low.'

----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK (skipped=1)

~roottwo\projects> python -m leveltest discover --level 5 -v
test_isupper (testproject.TestStringMethods) ... ok
test_split (testproject.TestStringMethods) ... ok
test_upper (testproject.TestStringMethods) ... ok

----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK

Or by specifying test cases to run, like:

~roottwo\projects>python -m leveltest --level 3 testproject -v
test_isupper (testproject.TestStringMethods) ... ok
test_split (testproject.TestStringMethods) ... skipped 'test level too low.'
test_upper (testproject.TestStringMethods) ... skipped 'test level too low.'

----------------------------------------------------------------------
Ran 3 tests in 0.002s

OK (skipped=2)

Answer 2

There is no way to pass arguments when using discover. DiscoveringTestLoader class from discover, removes all unmatched files (eliminates using '*test.py --level=3') and passes only file names into unittest.TextTestRunner

Probably only option so far is using environment variables

LEVEL=3 python -m unittest discoverfoo '*test.py'

Answer 3

The problem you have is that the unittest argument parser simply does not understand this syntax. You therefore have to remove the parameters before unittest is invoked.

A simple way to do this is to create a wrapper module (say my_unittest.py) that looks for your extra parameters, strips them from sys.argv and then invokes the main entry in unittest.

Now for the good bit... The code for that wrapper is basically the same as the code you already use for the single file case! You just need to put it into a separate file.

EDIT: Added sample code below as requested...

First, the new file to run the UTs (my_unittest.py):

import sys
import unittest
from parser import wrapper

if __name__ == '__main__':
    wrapper.parse_args()
    unittest.main(module=None, argv=sys.argv)

Now parser.py, which had to be in a separate file to avoid being in the __main__ module for the global reference to work:

import sys
import argparse
import unittest

class UnitTestParser(object):

    def __init__(self):
        self.args = None

    def parse_args(self):
        # Parse optional extra arguments
        parser = argparse.ArgumentParser()
        parser.add_argument('--level', type=int, default=0)
        ns, args = parser.parse_known_args()
        self.args = vars(ns)

        # Now set the sys.argv to the unittest_args (leaving sys.argv[0] alone)
        sys.argv[1:] = args

wrapper = UnitTestParser()

And finally a sample test case (project_test.py) to test that the parameters are parsed correctly:

import unittest
from parser import wrapper

class TestMyProject(unittest.TestCase):

    def test_len(self):
        self.assertEqual(len(wrapper.args), 1)

    def test_level3(self):
        self.assertEqual(wrapper.args['level'], 3)

And now the proof:

$ python -m my_unittest discover --level 3 . '*test.py'
..
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK