A module to make writing tests that deal with asynchronous things easier.

async_utils.py

import unittest

# A module to make writing tests that deal with asynchronous things easier.

def async(f):
  """Decorator to use on any test case that uses assertCalled or
  assertNotCalled."""
  def g(self):
    # run_after is a list of functions to run after the test case. We clear it
    # before and after running the test case so that AsyncTestCase can see that
    # we've used the decorator..
    self.run_after = []
    try:
      f(self)
      for func in self.run_after:
        func()
    finally:
      self.run_after = []
  return g


class AsyncTestCase(unittest.TestCase):
  """Add methods to make testing asynchronous things easier."""
  
  def __init__(self, *args, **kwargs):
    unittest.TestCase.__init__(self, *args, **kwargs)
    self.run_after = []
    
    # Check that run_after has been cleared by async.
    def check_run_after():
      assert not self.run_after, "Needs to be decorated with async."
    self.addCleanup(check_run_after)
  
  def assertCalled(self, msg="assertCalled"):
    """Return a function that must be called (with any arguments) for the test
    case to pass.
    """
    # The called flag needs to be mutable for closures to work right, so wrap
    # it in a list. This can be solved better with the nonlocal keyword in
    # python3.
    called = [False]
    def test_func(*args, **kwargs):
      called[0] = True
      
    def after():
      self.assertTrue(called[0], msg)
    
    self.run_after.append(after)
    return test_func
  
  def assertNotCalled(self, msg="assertNotCalled"):
    """Return a function that must not be called for the test case to pass."""
    # Same hack as above.
    called = [False]
    def test_func(*args, **kwargs):
      called[0] = True
      
    def after():
      self.assertFalse(called[0], msg)
    
    self.run_after.append(after)
    return test_func


class TestTest(AsyncTestCase):
  """Test that the extra methods added by AsyncTestCase work right.
  Very meta.
  """
  
  @async
  def test_assertCalled(self):
    self.assertCalled()()
  
  @async
  def test_assertNotCalled(self):
    self.assertNotCalled()
  
  @unittest.expectedFailure
  @async
  def test_assertCalled_fail(self):
    self.assertCalled()
  
  @unittest.expectedFailure
  @async
  def test_assertNotCalled_fail(self):
    self.assertNotCalled()()


if __name__ == '__main__':
  unittest.main()