def test_generate_code_works(upstream):
"""
The Wormhole.generate_code method should resolve into a code generated by
the upstream if all goes well.
"""
upstream.get_code = lambda: succeed('code')
wormhole = Wormhole()
res = yield wormhole.generate_code()
assert res == 'code'
def test_generate_code_timeout(upstream):
"""
The Wormhole.generate_code method should be able to handle timeouts.
"""
def get_code():
deferred = Deferred()
reactor.callLater(2, deferred.callback, 'code')
return deferred
upstream.get_code = get_code
wormhole = Wormhole()
with pytest.raises(Timeout):
yield wormhole.generate_code(timeout=0)