Added fail_on_none parameter for Client.consume function (#172)

This commit is contained in:
Krukov D 2016-05-29 04:31:15 +03:00 committed by Andrew Godwin
parent 917ba184bb
commit 1a09540ca8

View File

@ -101,21 +101,27 @@ class Client(object):
content.setdefault('reply_channel', self.reply_channel) content.setdefault('reply_channel', self.reply_channel)
self.channel_layer.send(to, content) self.channel_layer.send(to, content)
def consume(self, channel): def consume(self, channel, fail_on_none=True):
""" """
Get next message for channel name and run appointed consumer Get next message for channel name and run appointed consumer
""" """
message = self.get_next_message(channel) message = self.get_next_message(channel)
if message: if message:
consumer, kwargs = self.channel_layer.router.match(message) match = self.channel_layer.router.match(message)
return consumer(message, **kwargs) if match:
consumer, kwargs = match
return consumer(message, **kwargs)
elif fail_on_none:
raise AssertionError("Can't find consumer for message %s" % message)
elif fail_on_none:
raise AssertionError("No message for channel %s" % channel)
def send_and_consume(self, channel, content={}): def send_and_consume(self, channel, content={}, fail_on_none=True):
""" """
Reproduce full live cycle of the message Reproduce full live cycle of the message
""" """
self.send(channel, content) self.send(channel, content)
return self.consume(channel) return self.consume(channel, fail_on_none=fail_on_none)
def receive(self): def receive(self):
"""self.get_next_message(self.reply_channel) """self.get_next_message(self.reply_channel)