class TestSinkWorker(unittest.TestCase):
def setUp(self):
global executable
self.loop = asyncio.get_event_loop()
ex_h = self.loop.get_exception_handler()
def err_handler(context, error):
ex_h(context, error)
raise Exception(str(error))
self.loop.set_exception_handler(err_handler)
self.pool = ProcessPoolExecutor(max_workers=4)
self.publisher = Communicator()
self.sink_func = executable
self.pool.submit = MagicMock(return_value=_Future(1))
def test_sinking(self):
global callback
terminate_event = asyncio.Event()
sink_worker = SinkWorker(self.loop,
self.pool,
self.sink_func,
self.publisher,
asyncio.Event(),
terminate_event,
callback=callback)
future = self.loop.create_task(sink_worker.start())
self.loop.run_until_complete(fake_sleep(1))
terminate_event.set()
self.loop.create_task(self.publisher.consume(1))
self.loop.run_until_complete(fake_sleep(1))
self.pool.submit.assert_called_once_with(self.sink_func, None, 1)
assert future.exception() is None
class TestInOutWorker(unittest.TestCase):
def setUp(self):
global executable
self.loop = asyncio.get_event_loop()
self.pool = ProcessPoolExecutor(max_workers=4)
self.publisher = Communicator()
self.consumer = Communicator()
self.in_out_func = executable
self.consumer.consume = AsyncMock(return_value=None)
def test_sinking(self):
inout_worker = InOutWorker(self.loop, self.pool, self.in_out_func,
self.publisher, self.consumer,
asyncio.Event(), asyncio.Event())
self.loop.create_task(self.publisher.consume(1))
self.loop.create_task(inout_worker.start())
self.loop.run_until_complete(asyncio.sleep(.1))
self.consumer.consume.assert_called_once_with(1)
def tearDown(self):
for task in asyncio.Task.all_tasks():
task.cancel()