async def test_poll_observer(): main = ApplicationCommunicator(MainConsumer, { 'type': 'channel', 'channel': CHANNEL_MAIN }) await main.send_input({ 'type': TYPE_POLL, 'observer': 'test', 'interval': 2 }) channel_layer = get_channel_layer() # Nothing should be received in the first second. async with async_timeout.timeout(1): try: await channel_layer.receive(CHANNEL_WORKER) assert False except CancelledError: pass async with async_timeout.timeout(2): # Then after another second we should get a notification. notify = await channel_layer.receive(CHANNEL_WORKER) assert notify['type'] == TYPE_EVALUATE assert notify['observer'] == 'test'
async def test_folkrnn_consumer(): tune = RNNTune.objects.create(**FOLKRNN_IN) scope = {'type': 'channel', 'channel': 'folk_rnn'} communicator = ApplicationCommunicator(FolkRNNConsumer, scope) await communicator.send_input({ 'type': 'folkrnn.generate', 'id': tune.id, }) while RNNTune.objects.last().rnn_finished is None: await sleep(0.1) await communicator.send_input({'type': 'stop'}) await communicator.wait() with open(TUNE_PATH + f'/thesession_with_repeats_{tune.id}_raw') as f: assert f.read() == FOLKRNN_OUT_RAW correct_out = FOLKRNN_OUT\ .replace('X:1', f'X:{tune.id}')\ .replace('№1', f'№{tune.id}') with open(TUNE_PATH + f'/thesession_with_repeats_{tune.id}') as f: assert f.read() == correct_out tune = RNNTune.objects.last() assert tune.rnn_started is not None assert tune.rnn_started is not None assert tune.rnn_started < tune.rnn_finished assert tune.rnn_finished - tune.rnn_started < timedelta(seconds=5) assert tune.abc == correct_out
async def test_poll_observer(): poller = ApplicationCommunicator(PollObserversConsumer, { 'type': 'channel', 'channel': CHANNEL_POLL_OBSERVER, }) await poller.send_input({ 'type': TYPE_POLL_OBSERVER, 'observer': 'test', 'interval': 5, }) channel_layer = get_channel_layer() # Nothing should be received in the frist 4 seconds. async with async_timeout.timeout(4): try: await channel_layer.receive(CHANNEL_WORKER_NOTIFY) assert False except CancelledError: pass # Then after two more seconds we should get a notification. notify = await channel_layer.receive(CHANNEL_WORKER_NOTIFY) assert notify['type'] == TYPE_EVALUATE_OBSERVER assert notify['observer'] == 'test'
async def run_consumer(timeout=None): """Run the consumer until it finishes processing. :param timeout: Set maximum execution time before cancellation, or ``None`` (default) for unlimited. """ manager_channel = state.MANAGER_CONTROL_CHANNEL manager_scope = { "type": "control_event", "channel": manager_channel, } manager_app = ApplicationCommunicator(ManagerConsumer(), manager_scope) channel_layer = get_channel_layer() async def _consume_loop(channel, scope, app): """Run a loop to consume messages off the channels layer.""" message = await channel_layer.receive(channel) while message.get("type", {}) != "_resolwe_manager_quit": message.update(scope) await app.send_input(message) message = await channel_layer.receive(channel) consume_future = asyncio.ensure_future( _consume_loop(manager_channel, manager_scope, manager_app) ) with suppress(asyncio.TimeoutError): await asyncio.wait_for(consume_future, timeout=timeout) await manager_app.wait()
async def test_notify_by_version(event): communicator = ApplicationCommunicator(EventsWorker, {'type': 'test'}) event.title = 'abc' event.save() versions = Version.objects.get_for_object(event) version_id = versions[0].pk await communicator.send_input({ 'type': 'notify_slack_by_event_version', 'version_id': version_id, }) await communicator.wait(timeout=1)
async def test_on_welcome(self, mock_welcome): """ `irc.receive` called with a `welcome` command should call the `on_welcome` """ communicator = ApplicationCommunicator(AsyncIrcConsumer, {'type': 'irc'}) await communicator.send_input({ 'type': 'irc.receive', 'command': 'welcome', 'channel': 'test_channel', 'user': '******', 'body': None, }) await communicator.wait(timeout=.2) self.assertEqual(mock_welcome.call_count, 1)
async def test_send_message(self): """ `send_message` should format the correct message and return it to the server """ class SendMessageConsumer(AsyncIrcConsumer): async def test_message(self, event): await self.send_message('my_channel', 'Hello IRC!') communicator = ApplicationCommunicator(SendMessageConsumer, {'type': 'irc'}) await communicator.send_input({'type': 'test.message'}) event = await communicator.receive_output(timeout=1) self.assertEqual( event, { 'type': 'irc.send', 'command': 'message', 'channel': 'my_channel', 'body': 'Hello IRC!', })
async def run_consumer(timeout=None, dry_run=False): """Run the consumer until it finishes processing. :param timeout: Set maximum execution time before cancellation, or ``None`` (default) for unlimited. :param dry_run: If ``True``, don't actually dispatch messages, just dequeue them. Defaults to ``False``. """ channel = state.MANAGER_CONTROL_CHANNEL scope = { "type": "control_event", "channel": channel, } app = ApplicationCommunicator(ManagerConsumer, scope) channel_layer = get_channel_layer() async def _consume_loop(): """Run a loop to consume messages off the channels layer.""" while True: message = await channel_layer.receive(channel) if dry_run: continue if message.get("type", {}) == "_resolwe_manager_quit": break message.update(scope) await app.send_input(message) if timeout is None: await _consume_loop() try: # A further grace period to catch late messages. async with async_timeout.timeout(timeout or 1): await _consume_loop() except asyncio.TimeoutError: pass await app.wait()
async def test_send_command(self): """ `send_command` should format the correct message and return it to the server """ class SendCommandConsumer(AsyncIrcConsumer): async def test_command(self, event): await self.send_command('join', channel='my_channel') communicator = ApplicationCommunicator(SendCommandConsumer, {'type': 'irc'}) # Give the loop a beat to initialize the instance await communicator.send_input({'type': 'test.command'}) event = await communicator.receive_output(timeout=1) self.assertEqual( event, { 'type': 'irc.send', 'command': 'join', 'channel': 'my_channel', 'body': None, })
async def run_consumer(timeout=None): """Run the consumer until it finishes processing.""" channel = state.MANAGER_CONTROL_CHANNEL scope = { 'type': 'control_event', 'channel': channel, } from . import manager app = ApplicationCommunicator(ManagerConsumer, scope) channel_layer = get_channel_layer() async def _consume_loop(): """Run a loop to consume messages off the channels layer.""" while True: message = await channel_layer.receive(channel) if message.get('type', {}) == '_resolwe_manager_quit': break message.update(scope) await manager.sync_counter.inc('message-handling') await app.send_input(message) if timeout is None: await _consume_loop() try: # A further grace period to catch late messages. async with async_timeout.timeout(timeout or 1): await _consume_loop() except asyncio.TimeoutError: pass await app.wait() flush = getattr(channel_layer, 'flush', None) if flush: await flush()
async def test_worker_and_client(): client_consumer = URLRouter([url(r'^ws/(?P<subscriber_id>.+)$', ClientConsumer)]) client = WebsocketCommunicator(client_consumer, '/ws/test-session') worker = ApplicationCommunicator( WorkerConsumer, {'type': 'channel', 'channel': CHANNEL_WORKER_NOTIFY} ) # Connect client. status, _ = await client.connect() assert status is True # Create an observer. @database_sync_to_async def create_observer(): observer = QueryObserver( create_request(views.PaginatedViewSet, offset=0, limit=10) ) items = observer.evaluate() assert not items add_subscriber('test-session', observer.id) return observer.id observer_id = await create_observer() # Create a single model instance for the observer model. @database_sync_to_async def create_model(): return models.ExampleItem.objects.create(enabled=True, name="hello world").pk primary_key = await create_model() # Check that ORM signal was generated. channel_layer = get_channel_layer() notify = await channel_layer.receive(CHANNEL_WORKER_NOTIFY) assert notify['type'] == TYPE_ORM_NOTIFY_TABLE assert notify['kind'] == ORM_NOTIFY_KIND_CREATE assert notify['primary_key'] == str(primary_key) # Propagate notification to worker. await worker.send_input(notify) # Check that observer evaluation was requested. notify = await channel_layer.receive(CHANNEL_WORKER_NOTIFY) assert notify['type'] == TYPE_EVALUATE_OBSERVER assert notify['observer'] == observer_id # Propagate notification to worker. await worker.send_input(notify) response = await client.receive_json_from() assert response['msg'] == 'added' assert response['primary_key'] == 'id' assert response['order'] == 0 assert response['item'] == {'id': 1, 'enabled': True, 'name': 'hello world'} # No other messages should be sent. assert await client.receive_nothing() is True await client.disconnect() # Run observer again and it should remove itself because there are no more subscribers. await worker.send_input({'type': TYPE_EVALUATE_OBSERVER, 'observer': observer_id}) assert await worker.receive_nothing() is True # Ensure that subscriber and observer have been removed. @database_sync_to_async def check_subscribers(): assert observer_models.Subscriber.objects.all().count() == 0 assert observer_models.Observer.objects.all().count() == 0 await check_subscribers()
async def test_poll_observer_integration(): client_consumer = URLRouter([url(r'^ws/(?P<subscriber_id>.+)$', ClientConsumer)]) client = WebsocketCommunicator(client_consumer, '/ws/test-session') worker = ApplicationCommunicator( WorkerConsumer, {'type': 'channel', 'channel': CHANNEL_WORKER_NOTIFY} ) poller = ApplicationCommunicator( PollObserversConsumer, {'type': 'channel', 'channel': CHANNEL_POLL_OBSERVER} ) # Connect client. status, _ = await client.connect() assert status is True # Create an observer. @database_sync_to_async def create_observer(): observer = QueryObserver(create_request(views.PollingObservableViewSet)) items = observer.evaluate() assert len(items) == 1 add_subscriber('test-session', observer.id) return observer.id observer_id = await create_observer() # Ensure that a notification message was sent to the poller. channel_layer = get_channel_layer() notify = await channel_layer.receive(CHANNEL_POLL_OBSERVER) assert notify['type'] == TYPE_POLL_OBSERVER assert notify['interval'] == 5 assert notify['observer'] == observer_id # Dispatch notification to poller as our poller uses a dummy queue. await poller.send_input(notify) # Nothing should be received in the frist 4 seconds. async with async_timeout.timeout(4): try: await channel_layer.receive(CHANNEL_WORKER_NOTIFY) assert False except CancelledError: pass # Then after two more seconds we should get a notification. notify = await channel_layer.receive(CHANNEL_WORKER_NOTIFY) # Dispatch notification to worker. await worker.send_input(notify) # Ensure another notification message was sent to the poller. notify = await channel_layer.receive(CHANNEL_POLL_OBSERVER) assert notify['type'] == TYPE_POLL_OBSERVER assert notify['interval'] == 5 assert notify['observer'] == observer_id # Ensure client got notified of changes. response = await client.receive_json_from() assert response['msg'] == 'changed' assert response['primary_key'] == 'id' assert response['order'] == 0 assert response['item']['static'].startswith('This is a polling observable:') # No other messages should be sent. assert await client.receive_nothing() is True await client.disconnect()
async def test_throttle_observer(): with override_settings( DJANGO_REST_FRAMEWORK_REACTIVE={'throttle_rate': 2}): channel_layer = get_channel_layer() main = ApplicationCommunicator(MainConsumer, { 'type': 'channel', 'channel': CHANNEL_MAIN }) worker = ApplicationCommunicator(WorkerConsumer, { 'type': 'channel', 'channel': CHANNEL_WORKER }) @database_sync_to_async def create_observer(): observer = QueryObserver( create_request(views.ExampleItemViewSet, offset=0, limit=10)) items = observer.subscribe('test-session') return observer.id @database_sync_to_async def get_last_evaluation(observer_id): return observer_models.Observer.objects.get( id=observer_id).last_evaluation def throttle_count(observer_id, value=None): cache_key = throttle_cache_key(observer_id) if value is None: return cache.get(cache_key) else: return cache.set(cache_key, value) observer_id = await create_observer() # Test that observer is evaluated async with async_timeout.timeout(1): await worker.send_input({ 'type': TYPE_EVALUATE, 'observer': observer_id }) # Nothing should be in the main worker assert await main.receive_nothing() # Get last evaluation time for later comparisson last_evaluation = await get_last_evaluation(observer_id) # Throttle count should be 1 (one process) assert throttle_count(observer_id) == 1 # Ensure last_evaluated time change before the next test await asyncio.sleep(0.001) # Observer evaluation is delayed while another is evaluated await worker.send_input({ 'type': TYPE_EVALUATE, 'observer': observer_id }) # Delayed observer should be scheduled notify = await channel_layer.receive(CHANNEL_MAIN) assert notify['type'] == TYPE_POLL assert notify['observer'] == observer_id # Throttle count should be 2 (two processes) assert throttle_count(observer_id) == 2 # Observer should not be evaluated assert last_evaluation == await get_last_evaluation(observer_id) # Observer evaluation is discarded when delayed observer scheduled await worker.send_input({ 'type': TYPE_EVALUATE, 'observer': observer_id }) # Nothing should be in the main worker assert await main.receive_nothing() # Observer should not be evaluated assert last_evaluation == await get_last_evaluation(observer_id) # Throttle count should be 3 (three processes) assert throttle_count(observer_id) == 3
async def test_worker_and_client(): client_consumer = URLRouter( [path('ws/<slug:subscriber_id>', ClientConsumer().as_asgi())]) client = WebsocketCommunicator(client_consumer, '/ws/test-session') main = ApplicationCommunicator(MainConsumer(), { 'type': 'channel', 'channel': CHANNEL_MAIN }) worker = ApplicationCommunicator(WorkerConsumer(), { 'type': 'channel', 'channel': CHANNEL_WORKER }) # Connect client. status, _ = await client.connect() assert status is True # Create an observer. @database_sync_to_async def create_observer(): observer = QueryObserver( create_request(views.PaginatedViewSet, offset=0, limit=10)) items = observer.subscribe('test-session') assert not items return observer.id observer_id = await create_observer() await assert_subscribers(1) await assert_subscribers(1, observer_id) # Create a single model instance for the observer model. @database_sync_to_async def create_model(): return models.ExampleItem.objects.create(enabled=True, name="hello world").pk primary_key = await create_model() channel_layer = get_channel_layer() async with async_timeout.timeout(1): # Check that ORM signal was generated. notify = await channel_layer.receive(CHANNEL_MAIN) assert notify['type'] == TYPE_ORM_NOTIFY assert notify['kind'] == ORM_NOTIFY_KIND_CREATE assert notify['primary_key'] == str(primary_key) # Propagate notification to worker. await main.send_input(notify) # Check that observer evaluation was requested. notify = await channel_layer.receive(CHANNEL_WORKER) assert notify['type'] == TYPE_EVALUATE assert notify['observer'] == observer_id # Propagate notification to worker. await worker.send_input(notify) response = await client.receive_json_from() assert response['msg'] == 'added' assert response['primary_key'] == 'id' assert response['order'] == 0 assert response['item'] == { 'id': primary_key, 'enabled': True, 'name': 'hello world', } # No other messages should be sent. assert await client.receive_nothing() is True await client.disconnect() # Ensure that subscriber has been removed. await assert_subscribers(0) async with async_timeout.timeout(1): # Run observer again and it should skip evaluation because there are no more subscribers. await worker.send_input({ 'type': TYPE_EVALUATE, 'observer': observer_id }) assert await worker.receive_nothing() is True
async def test_poll_observer_integration(): client_consumer = URLRouter( [path('ws/<slug:subscriber_id>', ClientConsumer().as_asgi())]) client = WebsocketCommunicator(client_consumer, '/ws/test-session') main = ApplicationCommunicator(MainConsumer(), { 'type': 'channel', 'channel': CHANNEL_MAIN }) worker = ApplicationCommunicator(WorkerConsumer(), { 'type': 'channel', 'channel': CHANNEL_WORKER }) # Connect client. status, _ = await client.connect() assert status is True # Create an observer. @database_sync_to_async def create_observer(): observer = QueryObserver(create_request( views.PollingObservableViewSet)) items = observer.subscribe('test-session') assert len(items) == 1 return observer.id observer_id = await create_observer() await assert_subscribers(1) await assert_subscribers(1, observer_id) channel_layer = get_channel_layer() async with async_timeout.timeout(1): # Ensure that a notification message was sent to the poller. notify = await channel_layer.receive(CHANNEL_MAIN) assert notify['type'] == TYPE_POLL assert notify['interval'] == 2 assert notify['observer'] == observer_id # Dispatch notification to poller as our poller uses a dummy queue. await main.send_input(notify) # Nothing should be received in the first second. with pytest.raises(asyncio.TimeoutError): async with async_timeout.timeout(1): await channel_layer.receive(CHANNEL_WORKER) assert False async with async_timeout.timeout(2): # Then after another second we should get a notification. notify = await channel_layer.receive(CHANNEL_WORKER) # Dispatch notification to worker. await worker.send_input(notify) # Ensure another notification message was sent to the poller. notify = await channel_layer.receive(CHANNEL_MAIN) assert notify['type'] == TYPE_POLL assert notify['interval'] == 2 assert notify['observer'] == observer_id # Ensure client got notified of changes. response = await client.receive_json_from() assert response['msg'] == 'changed' assert response['primary_key'] == 'id' assert response['order'] == 0 assert response['item']['static'].startswith( 'This is a polling observable:') # No other messages should be sent. assert await client.receive_nothing() is True await client.disconnect()