def test_extraction_wrapper_for_object_with_call_method(self):
class FooBar(object):
def __call__(self, foo, bar, other='Hello'):
return (foo, bar, other)
FooBar.__call__ = extraction_wrapper(
FooBar.__call__, [
lambda env, context: dict(foo=context.get_foo()),
lambda env, context: dict(bar=env.get_bar())
],
function=False,
return_handlers=[lambda rval, context: context.add_value(rval)])
context = MagicMock()
context.get_foo = MagicMock(return_value='mock foo')
environment = MagicMock()
environment.get_bar = MagicMock(return_value='mock bar')
foo_bar = FooBar()
returned_value = foo_bar(environment, context)
context.get_foo.assert_called_once_with()
environment.get_bar.assert_called_once_with()
context.add_value.assert_called_once_with(
('mock foo', 'mock bar', 'Hello'))
self.assertEqual(returned_value, ('mock foo', 'mock bar', 'Hello'))
def test_extraction_wrapper_with_exception_handlers_unhandled_exception(
self):
def foo_bar_func(foo, bar, other='Hello'):
raise TypeError('Some exception message')
return (foo, bar, other)
wrapped_foo_bar_func = extraction_wrapper(
foo_bar_func, [
lambda env, context: dict(foo=context.get_foo()),
lambda env, context: dict(bar=env.get_bar())
],
exception_handlers=[(ValueError,
lambda e, context: 'Handled ValueError')])
context = MagicMock()
context.get_foo = MagicMock(return_value='mock foo')
environment = MagicMock()
environment.get_bar = MagicMock(return_value='mock bar')
with self.assertRaises(TypeError):
returned_value = wrapped_foo_bar_func(environment, context)
context.get_foo.assert_called_once_with()
environment.get_bar.assert_called_once_with()
self.assertEqual(context.add_value.call_count, 0)
def test_extraction_wrapper_for_function_with_default_extractor_and_no_return_handlers(
self):
def foo_bar_func(foo, bar, other='Hello'):
return (foo, bar, other)
wrapped_foo_bar_func = extraction_wrapper(foo_bar_func, [
lambda env, context: dict(foo=context.get_foo()),
lambda env, context: dict(bar=env.get_bar())
])
context = MagicMock()
context.get_foo = MagicMock(return_value='mock foo')
environment = MagicMock()
environment.get_bar = MagicMock(return_value='mock bar')
returned_value = wrapped_foo_bar_func(environment, context)
context.get_foo.assert_called_once_with()
environment.get_bar.assert_called_once_with()
self.assertEqual(context.add_value.call_count, 0)
self.assertEqual(returned_value, ('mock foo', 'mock bar', 'Hello'))
def test_extraction_wrapper_for_function_with_return_extractor(self):
def foo_bar_func(foo, bar, other='Hello'):
return (foo, bar, other)
wrapped_foo_bar_func = extraction_wrapper(
foo_bar_func, [
lambda env, context: dict(foo=context.get_foo()),
lambda env, context: dict(bar=env.get_bar())
],
return_handlers=[lambda rval, context: context.add_value(rval[2])],
return_value_extractor=lambda rval, context: rval[2])
context = MagicMock()
context.get_foo = MagicMock(return_value='mock foo')
environment = MagicMock()
environment.get_bar = MagicMock(return_value='mock bar')
returned_value = wrapped_foo_bar_func(environment, context)
context.get_foo.assert_called_once_with()
environment.get_bar.assert_called_once_with()
context.add_value.assert_called_once_with('Hello')
self.assertEqual(returned_value, 'Hello')
# We achieve this by a combination of create_conditional_step and extraction_wrapper.
# First we need to create an action function (that accepts TrailEnvironment, Context etc). This is done using
# extraction_wrapper.
be_lazy = create_conditional_step( # Make the Step conditional, i.e., skip progeny if the step fails.
# extraction_wrapper is a very powerful function conversion wrapper. Please look-up the full documentation
# for all the features it provides.
extraction_wrapper(
# The function that needs to be wrapped, i.e., converted into an action function. In this case it is:
lazy,
# Next we need to define all the parameter extractors, i.e., functions that will accept TrailEnvironment and
# Context and return the kwargs required by the wrapped function (lazy). Since we are extracting only one
# parameter and it can be expressed using a single lambda expression, we define our parameter extractors as
# the below list:
[
# This extracts the value from the context using the Context.get_branch_choice function and returns a
# dictionary.
lambda trail_env, context: dict(branch_choice=context.
get_branch_choice()),
]
# The dictionary produced as a result of all the parameter extractors is passed as **kwargs to the function
# lazy, in this case, it results in:
# lazy(branch_choice=context.get_branch_choice())
# By default, the return value of the function is the return value of the step. In this case, whatever string
# the function lazy produces, becomes the return value of the step.
))
# The above example is replicated in the step definition below.
work_hard = create_conditional_step(
extraction_wrapper(hardwork, [
lambda trail_env, context: dict(branch_choice=context.
get_branch_choice())
