Reproducible, random-data-driven tests for unittest

unittest-random.py

# -*- coding: utf-8 -*-
"""
    RandTest

    A unittest.TestCase subclass that facilitates working with
    randomly-generated test data.

    https://gist.github.com/2020459

    :copyright: (c) 2012 by Ori Livneh
    :license: MIT
"""
import base64
import os
import random
import unittest


__all__ = ('RandTest', 'main')


def b64seed(size=16):
    """ Generates a base64-encoded array of `size` bytes """
    return base64.b64encode(os.urandom(size))


class RandTestResult(unittest.TestResult):
    """ RandTest test result recorder """

    def __init__(self, *args, **kwargs):
        self.seed = kwargs.pop('seed', None)
        super(RandTestResult, self).__init__(*args, **kwargs)


class RandTest(unittest.TestCase):
    """ A unittest.TestCase for tests that use random data """

    def defaultTestResult(self):
        return RandTestResult()

    def run(self, result=None, seed=None):
        self.__seed = seed or getattr(result, 'seed', b64seed())
        random.seed(self.__seed)
        super(RandTest, self).run(result)

    def failureException(self, *args, **kwargs):
        args = list(args)
        args.append(self.__seed)
        return super(RandTest, self).failureException(*args, **kwargs)


class RandTestRunner(unittest.TextTestRunner):
    """
    A test runner that can be configured to seed the random number generator,
    for reproducible randomized tests.
    """
    resultclass = RandTestResult

    def __init__(self, *args, **kwargs):
        self.seed = kwargs.pop('seed', None)
        super(RandTestRunner, self).__init__(*args, **kwargs)

    def _makeResult(self):
        return self.resultclass(self.stream, self.descriptions,
                                self.verbosity, seed=self.seed)


class RandTestTestCase(RandTest):
    """ Tests for RandTest """

    def test_reproducibility(self):
        """ Random values should be reproducible given same seed """
        self.assertEqual(random.random(), 0.6799777634765871)

    def test_seed_generation(self):
        """ Seeds should decode to byte arrays of requested length """
        seed = base64.b64decode(b64seed(123))
        self.assertEqual(len(seed), 123)

    def test_failures(self):
        """ AssertionErrors should be annotated with the seed value """
        try:
            self.assertTrue(False)
        except AssertionError as ex:
            self.assertEquals(ex.args[1], 31337)
        else:
            self.fail()


def main(*args, **kwargs):
    """ Analog of unittest.main() """
    runner = RandTestRunner(
        seed=kwargs.pop('seed', None),
        verbosity=kwargs.get('verbosity', 1),
        buffer=kwargs.get('buffer', True)
    )

    kwargs.setdefault('testRunner', runner)
    unittest.main(*args, **kwargs)

if __name__ == '__main__':
    main(seed=31337, verbosity=2)