class TestPipeline(TestCase):
"""Tests for the main pipeline"""
def setUp(self):
self.pl = Pipeline()
def test_attach(self):
self.pl.attach(Module, 'module1')
self.pl.attach(Module, 'module2')
print([m.name for m in self.pl.modules])
self.assertEqual('module1', self.pl.modules[0].name)
self.assertEqual('module2', self.pl.modules[1].name)
def test_attach_bundle(self):
modules = [Module, Module]
self.pl.attach_bundle(modules)
self.assertEqual(2, len(self.pl.modules))
def test_attach_function(self):
self.pl.attach(lambda x: 1)
self.pl.attach(lambda x: 2, "Another Lambda")
self.assertEqual('<lambda>', self.pl.modules[0].name)
self.assertEqual('Another Lambda', self.pl.modules[1].name)
def test_drain_calls_each_attached_module(self):
pl = Pipeline(blob=1)
func_module_spy = MagicMock()
def func_module(blob):
func_module_spy()
return blob
pl.attach(Module, 'module1')
pl.attach(func_module, 'module2')
pl.attach(Module, 'module3')
for module in pl.modules:
print(module)
if isinstance(module, Module):
print("normal module, mocking")
module.process = MagicMock(return_value={})
n = 3
pl.drain(n)
for module in pl.modules:
try:
# Regular modules
self.assertEqual(n, module.process.call_count)
except AttributeError:
# Function module
self.assertEqual(n, func_module_spy.call_count)
def test_drain_calls_process_method_on_each_attached_module(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2')
pl.attach(Module, 'module3')
for module in pl.modules:
module.process = MagicMock(return_value={})
n = 3
pl.drain(n)
for module in pl.modules:
self.assertEqual(n, module.process.call_count)
def test_drain_doesnt_call_process_if_blob_is_none(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2')
pl.attach(Module, 'module3')
pl.modules[0].process = MagicMock(return_value=None)
pl.modules[1].process = MagicMock(return_value={})
pl.modules[2].process = MagicMock(return_value={})
n = 3
pl.drain(n)
self.assertEqual(n, pl.modules[0].process.call_count)
self.assertEqual(0, pl.modules[1].process.call_count)
self.assertEqual(0, pl.modules[2].process.call_count)
def test_conditional_module_not_called_if_key_not_in_blob(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2', only_if='foo')
pl.attach(Module, 'module3')
for module in pl.modules:
module.process = MagicMock(return_value={})
pl.drain(1)
self.assertEqual(1, pl.modules[0].process.call_count)
self.assertEqual(0, pl.modules[1].process.call_count)
self.assertEqual(1, pl.modules[2].process.call_count)
def test_conditional_module_not_called_if_multiple_keys_not_in_blob(self):
pl = Pipeline(blob=1)
to_be_called = MagicMock()
not_to_be_called = MagicMock()
class DummyPump(Pump):
def process(self, blob):
blob['foo'] = 1
blob['bar'] = 2
return blob
class ConditionalModule(Module):
def process(self, blob):
print(self.only_if)
print(blob)
not_to_be_called()
assert False
return blob
class Module1(Module):
def process(self, blob):
to_be_called()
return blob
class Module2(Module):
def process(self, blob):
to_be_called()
return blob
pl.attach(DummyPump)
pl.attach(Module1)
pl.attach(ConditionalModule, only_if=['foo', 'narf'])
pl.attach(ConditionalModule, only_if=['narf', 'bar'])
pl.attach(Module2)
pl.drain(3)
assert 6 == to_be_called.call_count
assert 0 == not_to_be_called.call_count
def test_conditional_module_called_if_key_in_blob(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2', only_if='foo')
pl.attach(Module, 'module3')
pl.modules[0].process = MagicMock(return_value={'foo': 23})
pl.modules[1].process = MagicMock(return_value={})
pl.modules[2].process = MagicMock(return_value={})
pl.drain(1)
self.assertEqual(1, pl.modules[0].process.call_count)
self.assertEqual(1, pl.modules[1].process.call_count)
self.assertEqual(1, pl.modules[2].process.call_count)
def test_conditional_module_called_if_multiple_keys_in_blob(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2', only_if=['foo', 'bar'])
pl.attach(Module, 'module3')
pl.modules[0].process = MagicMock(return_value={'foo': 23, 'bar': 5})
pl.modules[1].process = MagicMock(return_value={})
pl.modules[2].process = MagicMock(return_value={})
pl.drain(1)
self.assertEqual(1, pl.modules[0].process.call_count)
self.assertEqual(1, pl.modules[1].process.call_count)
self.assertEqual(1, pl.modules[2].process.call_count)
def test_condition_every(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2', every=3)
pl.attach(Module, 'module3', every=9)
pl.attach(Module, 'module4', every=10)
pl.attach(Module, 'module5')
func_module = MagicMock()
func_module.__name__ = "MagicMock"
pl.attach(func_module, 'funcmodule', every=4)
for module in pl.modules:
module.process = MagicMock(return_value={})
pl.drain(9)
self.assertEqual(9, pl.modules[0].process.call_count)
self.assertEqual(3, pl.modules[1].process.call_count)
self.assertEqual(1, pl.modules[2].process.call_count)
self.assertEqual(0, pl.modules[3].process.call_count)
self.assertEqual(9, pl.modules[4].process.call_count)
self.assertEqual(2, func_module.call_count)
def test_selective_blob_keys(self):
class DummyPump(Pump):
def process(self, blob):
return Blob({'a': 1, 'b': 2, 'c': 3})
class Observer(Module):
def configure(self):
self.needed_key = self.require('needed_key')
def process(self, blob):
print(blob)
assert 1 == len(blob)
assert self.needed_key in blob
return blob
pl = Pipeline()
pl.attach(DummyPump)
pl.attach(Observer, needed_key='a', blob_keys=['a'])
pl.attach(Observer, needed_key='b', blob_keys=['b'])
pl.attach(Observer, needed_key='c', blob_keys=['c'])
pl.drain(3)
def test_selective_blob_keys_with_multiple_keys(self):
n_cycles = 3
mock_to_be_called = MagicMock()
class DummyPump(Pump):
def process(self, blob):
return Blob({'a': 1, 'b': 2, 'c': 3})
class Observer(Module):
def process(self, blob):
assert 2 == len(blob)
assert 'a' in blob
assert 'b' in blob
mock_to_be_called()
return blob
class OtherObserver(Module):
def process(self, blob):
assert 3 == len(blob)
assert 'a' in blob
assert 'b' in blob
assert 'c' in blob
mock_to_be_called()
return blob
pl = Pipeline()
pl.attach(DummyPump)
pl.attach(Observer, blob_keys=['a', 'b'])
pl.attach(OtherObserver)
pl.drain(n_cycles)
assert 2 * n_cycles == mock_to_be_called.call_count
def test_selective_blob_keys_with_missing_key(self):
n_cycles = 3
mock_to_be_called = MagicMock()
class DummyPump(Pump):
def process(self, blob):
return Blob({'a': 1, 'b': 2, 'c': 3})
class Observer(Module):
def process(self, blob):
assert 0 == len(blob)
mock_to_be_called()
return blob
pl = Pipeline()
pl.attach(DummyPump)
pl.attach(Observer, blob_keys=['x'])
pl.drain(n_cycles)
assert n_cycles == mock_to_be_called.call_count
def test_selective_blob_keys_mutating_the_blob(self):
class DummyPump(Pump):
def process(self, blob):
return Blob({'a': 1, 'b': 2, 'c': 3})
class Mutator(Module):
def process(self, blob):
assert 1 == len(blob)
blob['d'] = 4
return blob
class Observer(Module):
def process(self, blob):
print(blob)
assert 4 == len(blob)
assert 'd' in blob
assert 4 == blob['d']
return blob
pl = Pipeline()
pl.attach(DummyPump)
pl.attach(Mutator, needed_key='a', blob_keys=['a'])
pl.attach(Observer)
pl.drain(3)
def test_selective_blob_keys_returning_nothing_doesnt_stop_the_cycle(self):
mock_to_be_called = MagicMock()
class DummyPump(Pump):
def process(self, blob):
return Blob({'a': 1, 'b': 2, 'c': 3})
class NoStopper(Module):
def process(self, blob):
return
class Observer(Module):
def process(self, blob):
mock_to_be_called()
assert 3 == len(blob)
return blob
pl = Pipeline()
pl.attach(DummyPump)
pl.attach(NoStopper, blob_keys=['a'])
pl.attach(Observer)
pl.drain(3)
assert 3 == mock_to_be_called.call_count
def test_drain_calls_function_modules(self):
pl = Pipeline(blob=1)
func_module1 = MagicMock()
func_module2 = MagicMock()
func_module3 = MagicMock()
func_module1.__name__ = "MagicMock"
func_module2.__name__ = "MagicMock"
func_module3.__name__ = "MagicMock"
pl.attach(func_module1, 'module1')
pl.attach(func_module2, 'module2')
pl.attach(func_module3, 'module3')
pl.drain(1)
self.assertEqual(1, pl.modules[0].call_count)
self.assertEqual(1, pl.modules[1].call_count)
self.assertEqual(1, pl.modules[2].call_count)
def test_finish(self):
out = self.pl.finish()
assert out is not None
def test_drain_calls_finish_on_each_attached_module(self):
self.pl.attach(Module, 'module1')
self.pl.attach(Module, 'module2')
self.pl.attach(lambda x: 1, 'func_module')
for module in self.pl.modules:
module.finish = MagicMock()
self.pl.drain(4)
for module in self.pl.modules:
if module.name != 'func_module':
self.assertEqual(1, module.finish.call_count)
def test_ctrl_c_handling(self):
pl = Pipeline()
self.assertFalse(pl._stop)
pl._handle_ctrl_c() # first KeyboardInterrupt
self.assertTrue(pl._stop)
with self.assertRaises(SystemExit):
pl._handle_ctrl_c() # second KeyboardInterrupt
def test_attaching_a_pump_allows_first_param_to_be_passed_as_fname(self):
class APump(Pump):
def configure(self):
self.filename = self.require('filename')
p = APump('test')
self.assertEqual('test', p.filename)
def test_attaching_multiple_pumps(self):
pl = Pipeline()
class APump(Pump):
def configure(self):
self.i = 0
def process(self, blob):
blob['A'] = self.i
self.i += 1
return blob
class BPump(Pump):
def configure(self):
self.i = 0
def process(self, blob):
blob['B'] = self.i
self.i += 1
return blob
class CheckBlob(Module):
def configure(self):
self.i = 0
def process(self, blob):
assert self.i == blob['A']
assert self.i == blob['B']
self.i += 1
return blob
pl.attach(APump)
pl.attach(BPump)
pl.attach(CheckBlob)
pl.drain(5)
def test_attached_module_gets_a_parameter_passed_which_is_ignored(self):
pl = Pipeline()
log_mock = MagicMock()
class A(Module):
def configure(self):
a = self.get('a')
self.log = log_mock
pl.attach(A, a=1, b=2)
pl.drain(1)
args, kwargs = log_mock.warning.call_args_list[0]
assert 'The following parameters were ignored: b' == args[0]
def test_attached_module_gets_multiple_parameters_passed_which_are_ignored(
self
):
pl = Pipeline()
log_mock = MagicMock()
class A(Module):
def configure(self):
a = self.get('a')
self.log = log_mock
pl.attach(A, a=1, b=2, c=3)
pl.drain(1)
args, kwargs = log_mock.warning.call_args_list[0]
assert 'The following parameters were ignored: b, c' == args[0]
def test_attached_module_does_not_warn_for_reserverd_parameters(self):
pl = Pipeline()
log_mock = MagicMock()
class A(Module):
def configure(self):
a = self.get('a')
self.log = log_mock
pl.attach(A, a=1, b=2, only_if='a', every=10)
pl.drain(1)
args, kwargs = log_mock.warning.call_args_list[0]
assert 'The following parameters were ignored: b' == args[0]
