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