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()
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]