def test_reading_multiple_corsika_files_pipeline(self):
class Tester(kp.Module):
def configure(self):
self.counter = 0
def process(self, blob):
self.counter = self.counter + 1
print(self.counter)
return blob
def finish(self):
assert (self.counter == 6)
pipe = Pipeline(timeit=False)
pipe.attach(
AanetPump,
filenames=[
join(
TEST_DATA_DIR,
'mcv5.0.DAT004340.propa.sirene.jte.jchain.aanet.4340.root'
),
join(
TEST_DATA_DIR,
'mcv5.0.DAT004340.propa.sirene.jte.jchain.aanet.4340.root'
)
]
)
pipe.attach(Tester)
pipe.drain()
def test_apply_to_hits_with_dom_id_and_channel_id(self):
hits = Table({
'dom_id': [2, 3, 3],
'channel_id': [0, 1, 2],
'time': [10.1, 11.2, 12.3]
})
tester = self
class HitCalibrator(Module):
def process(self, blob):
chits = self.services['calibrate'](hits)
assert len(hits) == len(chits)
a_hit = chits[0]
tester.assertAlmostEqual(2.1, a_hit.pos_x)
tester.assertAlmostEqual(40, a_hit.t0)
t0 = a_hit.t0
tester.assertAlmostEqual(10.1 + t0, a_hit.time)
a_hit = chits[1]
tester.assertAlmostEqual(3.4, a_hit.pos_x)
tester.assertAlmostEqual(80, a_hit.t0)
t0 = a_hit.t0
tester.assertAlmostEqual(11.2 + t0, a_hit.time)
return blob
pipe = Pipeline()
pipe.attach(CalibrationService, filename=DETX_FILENAME)
pipe.attach(HitCalibrator)
pipe.drain(1)
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')
self.assertEqual('module1', self.pl.modules[0].name)
self.assertEqual('module2', self.pl.modules[1].name)
def test_drain_calls_process_method_on_each_attached_module(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2')
for module in pl.modules:
module.process = MagicMock(return_value=1)
pl.drain(1)
for module in pl.modules:
module.process.assert_called_once_with(1)
def test_finish(self):
self.pl.finish()
def test_drain_calls_finish_on_each_attached_module(self):
self.pl.attach(Module, 'module1')
self.pl.attach(Module, 'module2')
for module in self.pl.modules:
module.finish = MagicMock()
self.pl.drain(4)
for module in self.pl.modules:
module.finish.assert_called_once_with()
def test_configuration_with_named_modules(self):
fobj = tempfile.NamedTemporaryFile(delete=True)
fobj.write(b"[X]\na = 1\nb = 2\n[Y]\nc='d'")
fobj.flush()
fname = str(fobj.name)
class A(Module):
def configure(self):
self.a = self.get('a')
self.b = self.get('b')
def process(self, blob):
assert 1 == self.a
assert 2 == self.b
return blob
class B(Module):
def configure(self):
self.c = self.get('c')
def process(self, blob):
assert 'd' == self.c
return blob
pipe = Pipeline(configfile=fname)
pipe.attach(A, 'X')
pipe.attach(B, 'Y')
pipe.drain(1)
fobj.close()
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_drain_calls_process_method_on_each_attached_module(self):
pl = Pipeline(blob=1)
pl.attach(Module, 'module1')
pl.attach(Module, 'module2')
for module in pl.modules:
module.process = MagicMock(return_value=1)
pl.drain(1)
for module in pl.modules:
module.process.assert_called_once_with(1)
def test_provided_detector_data(self):
class DetectorReader(Module):
def process(self, blob):
assert 'detector' in self.services
det = self.services['detector']
assert isinstance(det, Detector)
pipe = Pipeline()
pipe.attach(CalibrationService, filename=DETX_FILENAME)
pipe.attach(DetectorReader)
pipe.drain(1)
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_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]
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_parameter_with_differing_name(self):
fobj = tempfile.NamedTemporaryFile(delete=True)
fobj.write(b"[A]\na = 'abc'")
fobj.flush()
fname = str(fobj.name)
class A(Module):
def configure(self):
self.the_a = self.get('a')
def process(self, blob):
return 'abc' == self.the_a
pipe = Pipeline(configfile=fname)
pipe.attach(A)
pipe.drain(1)
fobj.close()
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_configuration_precedence_over_kwargs_when_require_is_used(self):
fobj = tempfile.NamedTemporaryFile(delete=True)
fobj.write(b"[A]\na = 1\n b = 'abc'")
fobj.flush()
fname = str(fobj.name)
class A(Module):
def configure(self):
self.xyz = self.require('a')
self.b = self.require('b')
def process(self, blob):
assert 1 == self.xyz
return 2 == self.b
pipe = Pipeline(configfile=fname)
pipe.attach(A)
pipe.drain(1)
fobj.close()
def test_configuration_with_config_for_a_module(self):
fobj = tempfile.NamedTemporaryFile(delete=True)
fobj.write(b"[A]\na = 1")
fobj.flush()
fname = str(fobj.name)
class A(Module):
def configure(self):
self.a = self.get('a')
def process(self, blob):
assert 1 == self.a
return blob
pipe = Pipeline(configfile=fname)
pipe.attach(A)
pipe.drain(1)
fobj.close()
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_configuration_precedence_over_kwargs(self):
fobj = tempfile.NamedTemporaryFile(delete=True)
fobj.write(b"[A]\na = 1\nb = 2")
fobj.flush()
fname = str(fobj.name)
class A(Module):
def configure(self):
self.a = self.get('a')
self.b = self.get('b')
def process(self, blob):
assert 1 == self.a
assert 2 == self.b
return blob
pipe = Pipeline(configfile=fname)
pipe.attach(A, b='foo')
pipe.drain(1)
fobj.close()
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_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_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_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_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_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_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_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 __init__(self, **context):
super(self.__class__, self).__init__(**context)
self.foo = self.get('foo') or 'default_foo'
self.bar = self.get('bar') or 23
def process(self, blob):
print("This is the current blob: " + str(blob))
blob['foo_entry'] = self.foo
return blob
class Moo(Module):
def process(self, blob):
blob['moo_entry'] = 42
return blob
class PrintBlob(Module):
def process(self, blob):
print(blob)
return blob
pipe = Pipeline()
pipe.attach(Pump, 'the_pump')
pipe.attach(Foo, 'foo_module', foo='dummyfoo', bar='dummybar')
pipe.attach(Moo, 'moo_module')
pipe.attach(PrintBlob, 'print_blob')
pipe.drain()
self.i = 0
def process(self, blob):
print("This is the current blob: " + str(blob))
self.i += 1
blob['foo_entry'] = self.foo
return blob
def finish(self):
print("My process() method was called {} times.".format(self.i))
def moo(blob):
"""A simple function to attach"""
blob['moo_entry'] = 42
return blob
class PrintBlob(Module):
def process(self, blob):
print(blob)
return blob
pipe = Pipeline()
pipe.attach(DummyPump, 'the_pump')
pipe.attach(Foo, bar='dummybar', baz=69)
pipe.attach(moo)
pipe.attach(PrintBlob)
pipe.drain()
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]
class TestServices(TestCase):
def setUp(self):
self.pl = Pipeline()
def test_service(self):
class Service(Module):
def configure(self):
self.expose(23, "foo")
self.expose(self.whatever, "whatever")
def whatever(self, x):
return x * 2
class UseService(Module):
def process(self, blob):
print(self.services)
assert 23 == self.services["foo"]
assert 2 == self.services["whatever"](1)
self.pl.attach(Service)
self.pl.attach(UseService)
self.pl.drain(1)
def test_service_usable_in_configure_when_attached_before(self):
return
class Service(Module):
def configure(self):
self.expose(23, "foo")
self.expose(self.whatever, "whatever")
def whatever(self, x):
return x * 2
class UseService(Module):
def configure(self):
assert 23 == self.services["foo"]
assert 2 == self.services["whatever"](1)
self.pl.attach(Service)
self.pl.attach(UseService)
self.pl.drain(1)
def test_required_service(self):
class AService(Module):
def configure(self):
self.expose(self.a_function, 'a_function')
def a_function(self, b='c'):
return b + 'd'
class AModule(Module):
def configure(self):
self.require_service('a_function', why='because')
def process(self, blob):
assert 'ed' == self.services['a_function']("e")
self.pl.attach(AService)
self.pl.attach(AModule)
self.pl.drain(2)
def test_required_service_not_present(self):
self.pl.log = MagicMock()
class AModule(Module):
def configure(self):
self.require_service('a_function', why='because')
def process(self, blob):
assert False # make sure that process is not called
self.pl.attach(AModule)
self.pl.drain(1)
self.pl.log.critical.assert_called_with(
'Following services are required and missing: a_function'
)
def test_required_service_not_present_in_multiple_modules(self):
self.pl.log = MagicMock()
class AModule(Module):
def configure(self):
self.require_service('a_function', why='because')
self.require_service('b_function', why='because')
def process(self, blob):
assert False # make sure that process is not called
class BModule(Module):
def configure(self):
self.require_service('c_function', why='because')
def process(self, blob):
assert False # make sure that process is not called
self.pl.attach(AModule)
self.pl.attach(BModule)
self.pl.drain(1)
self.pl.log.critical.assert_called_with(
'Following services are required and missing: '
'a_function, b_function, c_function'
)
def test_required_service_not_present_but_some_are_present(self):
self.pl.log = MagicMock()
class AModule(Module):
def configure(self):
self.expose(self.d_function, 'd_function')
self.require_service('a_function', why='because')
self.require_service('b_function', why='because')
def d_function(self):
pass
def process(self, blob):
assert False # make sure that process is not called
class BModule(Module):
def configure(self):
self.require_service('c_function', why='because')
def process(self, blob):
assert False # make sure that process is not called
class CModule(Module):
def configure(self):
self.require_service('d_function', why='because')
def process(self, blob):
assert False # make sure that process is not called
self.pl.attach(AModule)
self.pl.attach(BModule)
self.pl.attach(CModule)
self.pl.drain(1)
self.pl.log.critical.assert_called_with(
'Following services are required and missing: '
'a_function, b_function, c_function'
)