# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

"""
Tests for L{twisted.internet.defer.deferredGenerator} and related APIs.
"""

from __future__ import division, absolute_import

import sys

from twisted.internet import reactor

from twisted.trial import unittest

from twisted.internet.defer import waitForDeferred, deferredGenerator, Deferred
from twisted.internet.defer import inlineCallbacks, returnValue
from twisted.internet import defer


def getThing():
d = Deferred()
reactor.callLater(0, d.callback, "hi")
return d

def getOwie():
d = Deferred()
def CRAP():
d.errback(ZeroDivisionError('OMG'))
reactor.callLater(0, CRAP)
return d

# NOTE: most of the tests in DeferredGeneratorTests are duplicated
# with slightly different syntax for the InlineCallbacksTests below.

class TerminalException(Exception):
pass

class BaseDefgenTests:
"""
This class sets up a bunch of test cases which will test both
deferredGenerator and inlineCallbacks based generators. The subclasses
DeferredGeneratorTests and InlineCallbacksTests each provide the actual
generator implementations tested.
"""

def testBasics(self):
"""
Test that a normal deferredGenerator works. Tests yielding a
deferred which callbacks, as well as a deferred errbacks. Also
ensures returning a final value works.
"""

return self._genBasics().addCallback(self.assertEqual, 'WOOSH')

def testBuggy(self):
"""
Ensure that a buggy generator properly signals a Failure
condition on result deferred.
"""
return self.assertFailure(self._genBuggy(), ZeroDivisionError)

def testNothing(self):
"""Test that a generator which never yields results in None."""

return self._genNothing().addCallback(self.assertEqual, None)

def testHandledTerminalFailure(self):
"""
Create a Deferred Generator which yields a Deferred which fails and
handles the exception which results. Assert that the Deferred
Generator does not errback its Deferred.
"""
return self._genHandledTerminalFailure().addCallback(self.assertEqual, None)

def testHandledTerminalAsyncFailure(self):
"""
Just like testHandledTerminalFailure, only with a Deferred which fires
asynchronously with an error.
"""
d = defer.Deferred()
deferredGeneratorResultDeferred = self._genHandledTerminalAsyncFailure(d)
d.errback(TerminalException("Handled Terminal Failure"))
return deferredGeneratorResultDeferred.addCallback(
self.assertEqual, None)

def testStackUsage(self):
"""
Make sure we don't blow the stack when yielding immediately
available deferreds.
"""
return self._genStackUsage().addCallback(self.assertEqual, 0)

def testStackUsage2(self):
"""
Make sure we don't blow the stack when yielding immediately
available values.
"""
return self._genStackUsage2().addCallback(self.assertEqual, 0)




class DeferredGeneratorTests(BaseDefgenTests, unittest.TestCase):

# First provide all the generator impls necessary for BaseDefgenTests
def _genBasics(self):

x = waitForDeferred(getThing())
yield x
x = x.getResult()

self.assertEqual(x, "hi")

ow = waitForDeferred(getOwie())
yield ow
try:
ow.getResult()
except ZeroDivisionError as e:
self.assertEqual(str(e), 'OMG')
yield "WOOSH"
return
_genBasics = deferredGenerator(_genBasics)

def _genBuggy(self):
yield waitForDeferred(getThing())
1//0
_genBuggy = deferredGenerator(_genBuggy)


def _genNothing(self):
if 0: yield 1
_genNothing = deferredGenerator(_genNothing)

def _genHandledTerminalFailure(self):
x = waitForDeferred(defer.fail(TerminalException("Handled Terminal Failure")))
yield x
try:
x.getResult()
except TerminalException:
pass
_genHandledTerminalFailure = deferredGenerator(_genHandledTerminalFailure)


def _genHandledTerminalAsyncFailure(self, d):
x = waitForDeferred(d)
yield x
try:
x.getResult()
except TerminalException:
pass
_genHandledTerminalAsyncFailure = deferredGenerator(_genHandledTerminalAsyncFailure)


def _genStackUsage(self):
for x in range(5000):
# Test with yielding a deferred
x = waitForDeferred(defer.succeed(1))
yield x
x = x.getResult()
yield 0
_genStackUsage = deferredGenerator(_genStackUsage)

def _genStackUsage2(self):
for x in range(5000):
# Test with yielding a random value
yield 1
yield 0
_genStackUsage2 = deferredGenerator(_genStackUsage2)

# Tests unique to deferredGenerator

def testDeferredYielding(self):
"""
Ensure that yielding a Deferred directly is trapped as an
error.
"""
# See the comment _deferGenerator about d.callback(Deferred).
def _genDeferred():
yield getThing()
_genDeferred = deferredGenerator(_genDeferred)

return self.assertFailure(_genDeferred(), TypeError)



class InlineCallbacksTests(BaseDefgenTests, unittest.TestCase):
# First provide all the generator impls necessary for BaseDefgenTests

def _genBasics(self):

x = yield getThing()

self.assertEqual(x, "hi")

try:
ow = yield getOwie()
except ZeroDivisionError as e:
self.assertEqual(str(e), 'OMG')
returnValue("WOOSH")
_genBasics = inlineCallbacks(_genBasics)

def _genBuggy(self):
yield getThing()
1/0
_genBuggy = inlineCallbacks(_genBuggy)


def _genNothing(self):
if 0: yield 1
_genNothing = inlineCallbacks(_genNothing)


def _genHandledTerminalFailure(self):
try:
x = yield defer.fail(TerminalException("Handled Terminal Failure"))
except TerminalException:
pass
_genHandledTerminalFailure = inlineCallbacks(_genHandledTerminalFailure)


def _genHandledTerminalAsyncFailure(self, d):
try:
x = yield d
except TerminalException:
pass
_genHandledTerminalAsyncFailure = inlineCallbacks(
_genHandledTerminalAsyncFailure)


def _genStackUsage(self):
for x in range(5000):
# Test with yielding a deferred
x = yield defer.succeed(1)
returnValue(0)
_genStackUsage = inlineCallbacks(_genStackUsage)

def _genStackUsage2(self):
for x in range(5000):
# Test with yielding a random value
yield 1
returnValue(0)
_genStackUsage2 = inlineCallbacks(_genStackUsage2)

# Tests unique to inlineCallbacks

def testYieldNonDeferrred(self):
"""
Ensure that yielding a non-deferred passes it back as the
result of the yield expression.
"""
def _test():
x = yield 5
returnValue(5)
_test = inlineCallbacks(_test)

return _test().addCallback(self.assertEqual, 5)

def testReturnNoValue(self):
"""Ensure a standard python return results in a None result."""
def _noReturn():
yield 5
return
_noReturn = inlineCallbacks(_noReturn)

return _noReturn().addCallback(self.assertEqual, None)

def testReturnValue(self):
"""Ensure that returnValue works."""
def _return():
yield 5
returnValue(6)
_return = inlineCallbacks(_return)

return _return().addCallback(self.assertEqual, 6)


def test_nonGeneratorReturn(self):
"""
Ensure that C{TypeError} with a message about L{inlineCallbacks} is
raised when a non-generator returns something other than a generator.
"""
def _noYield():
return 5
_noYield = inlineCallbacks(_noYield)

self.assertIn("inlineCallbacks",
str(self.assertRaises(TypeError, _noYield)))


def test_nonGeneratorReturnValue(self):
"""
Ensure that C{TypeError} with a message about L{inlineCallbacks} is
raised when a non-generator calls L{returnValue}.
"""
def _noYield():
returnValue(5)
_noYield = inlineCallbacks(_noYield)

self.assertIn("inlineCallbacks",
str(self.assertRaises(TypeError, _noYield)))