twistedtools.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. """
  2. Twisted integration
  3. -------------------
  4. This module provides a very simple way to integrate your tests with the
  5. Twisted_ event loop.
  6. You must import this module *before* importing anything from Twisted itself!
  7. Example::
  8. from nose.twistedtools import reactor, deferred
  9. @deferred()
  10. def test_resolve():
  11. return reactor.resolve("www.python.org")
  12. Or, more realistically::
  13. @deferred(timeout=5.0)
  14. def test_resolve():
  15. d = reactor.resolve("www.python.org")
  16. def check_ip(ip):
  17. assert ip == "67.15.36.43"
  18. d.addCallback(check_ip)
  19. return d
  20. .. _Twisted: http://twistedmatrix.com/trac/
  21. """
  22. import sys
  23. from Queue import Queue, Empty
  24. from nose.tools import make_decorator, TimeExpired
  25. __all__ = [
  26. 'threaded_reactor', 'reactor', 'deferred', 'TimeExpired',
  27. 'stop_reactor'
  28. ]
  29. _twisted_thread = None
  30. def threaded_reactor():
  31. """
  32. Start the Twisted reactor in a separate thread, if not already done.
  33. Returns the reactor.
  34. The thread will automatically be destroyed when all the tests are done.
  35. """
  36. global _twisted_thread
  37. try:
  38. from twisted.internet import reactor
  39. except ImportError:
  40. return None, None
  41. if not _twisted_thread:
  42. from twisted.python import threadable
  43. from threading import Thread
  44. _twisted_thread = Thread(target=lambda: reactor.run( \
  45. installSignalHandlers=False))
  46. _twisted_thread.setDaemon(True)
  47. _twisted_thread.start()
  48. return reactor, _twisted_thread
  49. # Export global reactor variable, as Twisted does
  50. reactor, reactor_thread = threaded_reactor()
  51. def stop_reactor():
  52. """Stop the reactor and join the reactor thread until it stops.
  53. Call this function in teardown at the module or package level to
  54. reset the twisted system after your tests. You *must* do this if
  55. you mix tests using these tools and tests using twisted.trial.
  56. """
  57. global _twisted_thread
  58. def stop_reactor():
  59. '''Helper for calling stop from withing the thread.'''
  60. reactor.stop()
  61. reactor.callFromThread(stop_reactor)
  62. reactor_thread.join()
  63. for p in reactor.getDelayedCalls():
  64. if p.active():
  65. p.cancel()
  66. _twisted_thread = None
  67. def deferred(timeout=None):
  68. """
  69. By wrapping a test function with this decorator, you can return a
  70. twisted Deferred and the test will wait for the deferred to be triggered.
  71. The whole test function will run inside the Twisted event loop.
  72. The optional timeout parameter specifies the maximum duration of the test.
  73. The difference with timed() is that timed() will still wait for the test
  74. to end, while deferred() will stop the test when its timeout has expired.
  75. The latter is more desireable when dealing with network tests, because
  76. the result may actually never arrive.
  77. If the callback is triggered, the test has passed.
  78. If the errback is triggered or the timeout expires, the test has failed.
  79. Example::
  80. @deferred(timeout=5.0)
  81. def test_resolve():
  82. return reactor.resolve("www.python.org")
  83. Attention! If you combine this decorator with other decorators (like
  84. "raises"), deferred() must be called *first*!
  85. In other words, this is good::
  86. @raises(DNSLookupError)
  87. @deferred()
  88. def test_error():
  89. return reactor.resolve("xxxjhjhj.biz")
  90. and this is bad::
  91. @deferred()
  92. @raises(DNSLookupError)
  93. def test_error():
  94. return reactor.resolve("xxxjhjhj.biz")
  95. """
  96. reactor, reactor_thread = threaded_reactor()
  97. if reactor is None:
  98. raise ImportError("twisted is not available or could not be imported")
  99. # Check for common syntax mistake
  100. # (otherwise, tests can be silently ignored
  101. # if one writes "@deferred" instead of "@deferred()")
  102. try:
  103. timeout is None or timeout + 0
  104. except TypeError:
  105. raise TypeError("'timeout' argument must be a number or None")
  106. def decorate(func):
  107. def wrapper(*args, **kargs):
  108. q = Queue()
  109. def callback(value):
  110. q.put(None)
  111. def errback(failure):
  112. # Retrieve and save full exception info
  113. try:
  114. failure.raiseException()
  115. except:
  116. q.put(sys.exc_info())
  117. def g():
  118. try:
  119. d = func(*args, **kargs)
  120. try:
  121. d.addCallbacks(callback, errback)
  122. # Check for a common mistake and display a nice error
  123. # message
  124. except AttributeError:
  125. raise TypeError("you must return a twisted Deferred "
  126. "from your test case!")
  127. # Catch exceptions raised in the test body (from the
  128. # Twisted thread)
  129. except:
  130. q.put(sys.exc_info())
  131. reactor.callFromThread(g)
  132. try:
  133. error = q.get(timeout=timeout)
  134. except Empty:
  135. raise TimeExpired("timeout expired before end of test (%f s.)"
  136. % timeout)
  137. # Re-raise all exceptions
  138. if error is not None:
  139. exc_type, exc_value, tb = error
  140. raise exc_type, exc_value, tb
  141. wrapper = make_decorator(func)(wrapper)
  142. return wrapper
  143. return decorate