def test_as_routes(self):
class Test(Consumers):
channel_name = 'test'
@consumer(tag='test')
def test(this, message):
return message
@consumer
def do(this, message):
return 'do'
def on_connect(this, message):
return 'connect'
channel_layer = channel_layers[DEFAULT_CHANNEL_LAYER]
routes = Test.as_routes(channel_name='test', path='^new$')
self.assertTrue(isinstance(routes, include))
self.assertEqual(routes.channel_names(), {'websocket.receive', 'websocket.connect',
'test', 'websocket.disconnect'})
self.assertEqual(len(routes.routing), 5)
self.assertEqual(routes.match(Message({'new': ''}, 'websocket.receive', channel_layer)), None)
m = Message({'path': 'new'}, 'websocket.connect', channel_layer)
self.assertEqual(routes.match(m)[0](m), 'connect')
message = Message({'path': 'new', 'tag': 'test'}, 'websocket.receive', channel_layer)
routes.match(message)[0](message)
self.assertEqual(channel_layer.receive_many(['test', ])[1], {'path': 'new', 'tag': 'test'})
def factory(name, **content):
channel_layer = ChannelLayer()
message = Message(content, name, channel_layer)
settings.SESSION_FILE_PATH = str(tmpdir)
message.channel_session = SessionStore()
return message
def test_enforce_ordering_slight(self):
"""
Tests that slight mode of enforce_ordering works
"""
# Construct messages to send
message0 = Message({
"reply_channel": "test-reply-a",
"order": 0
}, "websocket.connect", channel_layers[DEFAULT_CHANNEL_LAYER])
message1 = Message({
"reply_channel": "test-reply-a",
"order": 1
}, "websocket.receive", channel_layers[DEFAULT_CHANNEL_LAYER])
message2 = Message({
"reply_channel": "test-reply-a",
"order": 2
}, "websocket.receive", channel_layers[DEFAULT_CHANNEL_LAYER])
# Run them in an acceptable slight order
@enforce_ordering(slight=True)
def inner(message):
pass
inner(message0)
inner(message2)
inner(message1)
# Ensure wait channel is empty
wait_channel = "__wait__.%s" % "test-reply-a"
next_message = self.get_next_message(wait_channel)
self.assertEqual(next_message, None)
def test_channel_and_http_session(self):
"""
Tests that channel_and_http_session decorator stores the http session key and hydrates it when expected
"""
# Make a session to try against
session = session_for_reply_channel("test-reply-session")
# Construct message to send
message = Message({
"reply_channel": "test-reply-session",
"http_version": "1.1",
"method": "GET",
"path": "/test2/",
"headers": {
"host": b"example.com",
"cookie": ("%s=%s" % (settings.SESSION_COOKIE_NAME, session.session_key)).encode("ascii"),
},
}, None, None)
@channel_and_http_session
def inner(message):
pass
inner(message)
# It should store the session key
self.assertEqual(message.channel_session[settings.SESSION_COOKIE_NAME], session.session_key)
# Construct a new message
message2 = Message({"reply_channel": "test-reply-session", "path": "/"}, None, None)
inner(message2)
# It should hydrate the http_session
self.assertEqual(message2.http_session.session_key, session.session_key)
def test_enforce_ordering_fail(self):
"""
Tests that strict mode of enforce_ordering fails on bad ordering
"""
# Construct messages to send
message0 = Message({
"reply_channel": "test-reply-c",
"order": 0
}, "websocket.connect", channel_layers[DEFAULT_CHANNEL_LAYER])
message2 = Message({
"reply_channel": "test-reply-c",
"order": 2
}, "websocket.receive", channel_layers[DEFAULT_CHANNEL_LAYER])
# Run them in an acceptable strict order
@enforce_ordering
def inner(message):
pass
inner(message0)
inner(message2)
# Ensure wait channel is not empty
wait_channel = "__wait__.%s" % "test-reply-c"
next_message = self.get_next_message(wait_channel)
self.assertNotEqual(next_message, None)
def test_enforce_ordering_concurrent(self):
"""
Tests that strict mode of enforce_ordering puts messages in the correct queue after
the current message number changes while the message is being processed
"""
# Construct messages to send
message0 = Message(
{"reply_channel": "test-reply-e", "order": 0},
"websocket.connect",
channel_layers[DEFAULT_CHANNEL_LAYER]
)
message2 = Message(
{"reply_channel": "test-reply-e", "order": 2},
"websocket.receive",
channel_layers[DEFAULT_CHANNEL_LAYER]
)
message3 = Message(
{"reply_channel": "test-reply-e", "order": 3},
"websocket.receive",
channel_layers[DEFAULT_CHANNEL_LAYER]
)
@channel_session
def add_session(message):
pass
# Run them in an acceptable strict order
@enforce_ordering
def inner(message):
pass
inner(message0)
inner(message3)
# Add the session now so it can be mocked
add_session(message2)
with mock.patch.object(message2.channel_session, 'load', return_value={'__channels_next_order': 2}):
inner(message2)
# Ensure wait channel is empty
wait_channel = "__wait__.%s" % "test-reply-e"
next_message = self.get_next_message(wait_channel)
self.assertEqual(next_message, None)
# Ensure messages 3 and 2 both ended up back on the original channel
expected = {
2: message2,
3: message3
}
for m in range(2):
message = self.get_next_message("websocket.receive")
expected.pop(message.content['order'])
self.assertEqual(expected, {})
def test_http_session(self):
"""
Tests that http_session correctly extracts a session cookie.
"""
# Make a session to try against
session1 = session_for_reply_channel("test-reply")
# Construct message to send
message = Message(
{
"reply_channel": "test-reply",
"http_version": "1.1",
"method": "GET",
"path": "/test2/",
"headers": {
"host":
b"example.com",
"cookie":
("%s=%s" % (settings.SESSION_COOKIE_NAME,
session1.session_key)).encode("ascii"),
},
}, None, None)
# Run it through http_session, make sure it works (test double here too)
@http_session
@http_session
def inner(message):
message.http_session["species"] = "horse"
inner(message)
# Check value assignment stuck
session2 = session_for_reply_channel("test-reply")
self.assertEqual(session2["species"], "horse")
def test_as_consumers_with_decor(self):
def decor(_consumer):
@wraps(_consumer)
def _wrap(message, *args, **kwargs):
message.decor = True
return _consumer(message, *args, **kwargs)
return _wrap
def decor2(_consumer):
@wraps(_consumer)
def _wrap(message, *args, **kwargs):
message.decor2 = True
return _consumer(message, *args, **kwargs)
return _wrap
class Test(Consumers):
path = '^/(?P<slug>[^/]+)'
channel_name = 'test'
decorators = [decor, ]
slug = 'slug'
def on_connect(this, message, slug=None):
this.slug = slug
return message
@consumer(tag='(?P<tag>[^/]+)', decorators=[decor2])
def tags(this, message, tag):
return this.message, this.kwargs, this.slug
routes = Test.as_routes()
channel_layer = channel_layers[DEFAULT_CHANNEL_LAYER]
message = Message({'path': '/new'}, 'websocket.connect', channel_layer)
_consumer, kwargs = routes.match(message)
self.assertEqual(kwargs, {'slug': 'new'})
self.assertEqual(_consumer.__name__, 'ws_connect')
self.assertTrue(_consumer(message, **kwargs).decor)
message = Message({'path': '/new', 'tag': 'test'}, 'test', channel_layer)
_consumer, kwargs = routes.match(message)
self.assertEqual(kwargs, {'tag': 'test'})
self.assertTrue(_consumer(message, **kwargs)[0].decor)
self.assertTrue(_consumer(message, **kwargs)[0].decor2)
self.assertEqual(_consumer(message, **kwargs)[1], {'tag': 'test'})
self.assertEqual(_consumer(message, **kwargs)[2], 'slug')
def test_invalid_origin_header(self):
invalid_headers = [
[], # origin header missing
[b'origin', b''], # origin header empty
[b'origin', b'\xc3\xa4'] # non-ascii
]
for headers in invalid_headers:
content = {'headers': [headers]}
message = Message(content, 'websocket.connect', None)
self.assertRaises(DenyConnection, connect, message)
def run(self):
"""
Tries to continually dispatch messages to consumers.
"""
if self.signal_handlers:
self.install_signal_handler()
channels = self.apply_channel_filters(
self.channel_layer.router.channels)
logger.info("Listening on channels %s", ", ".join(sorted(channels)))
while not self.termed:
self.in_job = False
channel, content = self.channel_layer.receive(channels, block=True)
self.in_job = True
# If no message, stall a little to avoid busy-looping then continue
if channel is None:
print "no_data"
time.sleep(0.01)
continue
# Create message wrapper
logger.debug("Got message on %s (reply %s)", channel,
content.get("reply_channel", "none"))
message = Message(
content=content,
channel_name=channel,
channel_layer=self.channel_layer,
)
# Add attribute to message if it's been retried almost too many times,
# and would be thrown away this time if it's requeued. Used for helpful
# warnings in decorators and such - don't rely on this as public API.
if content.get("__retries__", 0) == self.message_retries:
message.__doomed__ = True
# Handle the message
match = self.channel_layer.router.match(message)
if match is None:
logger.error(
"Could not find match for message on %s! Check your routing.",
channel)
continue
gevent.spawn(Worker.foo, self, channel, content, message, match)
def test_as_routes_without_custom_routes(self):
class Test(Consumers):
path = '^new$'
channel_name = 'test'
def on_connect(this, message):
return 'connect'
routes = Test.as_routes()
self.assertTrue(isinstance(routes, include))
self.assertEqual(routes.channel_names(), {'websocket.receive', 'websocket.connect', 'websocket.disconnect'})
self.assertEqual(len(routes.routing), 3)
channel_layer = channel_layers[DEFAULT_CHANNEL_LAYER]
self.assertEqual(routes.match(Message({'new': ''}, 'websocket.receive', channel_layer)), None)
m = Message({'path': 'new'}, 'websocket.connect', channel_layer)
self.assertEqual(routes.match(m)[0](m), 'connect')
message = Message({'path': 'new', 'tag': 'test'}, 'websocket.receive', channel_layer)
self.assertEqual(routes.match(message)[0](message), None)
def ws_connect(message: Message):
# Accept connection
message.reply_channel.send({'accept': True})
# Work out room name from path (ignore slashes)
room = message.content['path'].strip('/')
# Save room in session and add us to the group
message.channel_session['room'] = room
# Add them to channel group
Group(f'chat-{room}').add(message.reply_channel)
def test_enforce_ordering_fail_no_order(self):
"""
Makes sure messages with no "order" key fail
"""
message0 = Message({"reply_channel": "test-reply-d"}, None,
channel_layers[DEFAULT_CHANNEL_LAYER])
@enforce_ordering
def inner(message):
pass
with self.assertRaises(ValueError):
inner(message0)
def test_channel_session_no_reply(self):
"""
Tests the channel_session decorator detects no reply channel
"""
# Construct message to send
message = Message({}, None, None)
# Run through a simple fake consumer that should trigger the error
@channel_session
@channel_session
def inner(message):
message.channel_session["num_ponies"] = -1
with self.assertRaises(ValueError):
inner(message)
def test_channel_session_third_arg(self):
"""
Tests the channel_session decorator with message as 3rd argument
"""
# Construct message to send
message = Message({"reply_channel": "test-reply"}, None, None)
# Run through a simple fake consumer that assigns to it
@channel_session
def inner(a, b, message):
message.channel_session["num_ponies"] = -1
with self.assertRaisesMessage(
ValueError, 'channel_session called without Message instance'):
inner(None, None, message)
def test_channel_session(self):
"""
Tests the channel_session decorator
"""
# Construct message to send
message = Message({"reply_channel": "test-reply"}, None, None)
# Run through a simple fake consumer that assigns to it
@channel_session
def inner(message):
message.channel_session["num_ponies"] = -1
inner(message)
# Test the session worked
session2 = session_for_reply_channel("test-reply")
self.assertEqual(session2["num_ponies"], -1)
def test_channel_session_no_reply_third_arg(self):
"""
Tests the channel_session decorator detects no reply channel
"""
# Construct message to send
message = Message({}, None, None)
# Run through a simple fake consumer that should trigger the error
@channel_session
@channel_session
def inner(a, b, message):
message.channel_session["num_ponies"] = -1
with self.assertRaisesMessage(
ValueError, 'channel_session called without Message instance'):
inner(None, None, message)
def test_channel_session_double_third_arg(self):
"""
Tests the channel_session decorator detects being wrapped in itself
and doesn't blow up.
"""
# Construct message to send
message = Message({"reply_channel": "test-reply"}, None, None)
# Run through a simple fake consumer that should trigger the error
@channel_session
@channel_session
def inner(a, b, message):
message.channel_session["num_ponies"] = -1
with self.assertRaisesMessage(
ValueError, 'channel_session called without Message instance'):
inner(None, None, message)
def get_next_message(self,
channel,
alias=DEFAULT_CHANNEL_LAYER,
require=False):
"""
Gets the next message that was sent to the channel during the test,
or None if no message is available.
If require is true, will fail the test if no message is received.
"""
recv_channel, content = channel_layers[alias].receive_many([channel])
if recv_channel is None:
if require:
self.fail("Expected a message on channel %s, got none" %
channel)
else:
return None
return Message(content, recv_channel, channel_layers[alias])
def test_channel_session_double(self):
"""
Tests the channel_session decorator detects being wrapped in itself
and doesn't blow up.
"""
# Construct message to send
message = Message({"reply_channel": "test-reply"}, None, None)
# Run through a simple fake consumer that should trigger the error
@channel_session
@channel_session
def inner(message):
message.channel_session["num_ponies"] = -1
inner(message)
# Test the session worked
session2 = session_for_reply_channel("test-reply")
self.assertEqual(session2["num_ponies"], -1)
def assertRoute(self, router, channel, content, consumer, kwargs=None):
"""
Asserts that asking the `router` to route the `content` as a message
from `channel` means it returns consumer `consumer`, optionally
testing it also returns `kwargs` to be passed in
Use `consumer` = None to assert that no route is found.
"""
message = Message(content, channel, channel_layer="fake channel layer")
match = router.match(message)
if match is None:
if consumer is None:
return
else:
self.fail("No route found for %s on %s; expecting %s" % (
content,
channel,
name_that_thing(consumer),
))
else:
mconsumer, mkwargs = match
if consumer is None:
self.fail("Route found for %s on %s; expecting no route." % (
content,
channel,
))
self.assertEqual(
consumer, mconsumer,
"Route found for %s on %s; but wrong consumer (%s not %s)." % (
content,
channel,
name_that_thing(mconsumer),
name_that_thing(consumer),
))
if kwargs is not None:
self.assertEqual(
kwargs, mkwargs,
"Route found for %s on %s; but wrong kwargs (%s not %s)." %
(
content,
channel,
mkwargs,
kwargs,
))
def test_invalid_origin(self):
content = {'headers': [[b'origin', b'http://example.org']]}
message = Message(content, 'websocket.connect', None)
self.assertRaises(DenyConnection, connect, message)
def test_valid_origin(self):
content = {'headers': [[b'origin', b'http://example.com']]}
message = Message(content, 'websocket.connect', None)
self.assertTrue(connect(message))
def factory(name, **content):
channel_layer = InMemoryChannelLayer()
message = Message(content, name, channel_layer)
settings.SESSION_FILE_PATH = str(tmpdir)
message.channel_session = FileSessionStore()
return message
def test_ws_message(self, engine):
message = Message({'reply_channel': 'test-reply'}, None, None)
ws_message(message)
engine.dispatch.assert_called_with(message)
def setUp(self):
"""
Make an in memory channel layer for testing
"""
self.channel_layer = ChannelLayer()
self.make_message = lambda m, c: Message(m, c, self.channel_layer)
def test_ws_disconnect(self, engine):
message = Message({'reply_channel': 'test-reply'}, None, None)
ws_disconnect(message)
engine.assert_called_with(message)
engine.disconnect.assert_called()
