I recently updated my project to Django 2 and channels 2. Right now I am trying to rewrite my tests for chat app.
I am facing a problem with tests that depend on django db mark from pytest-django. I tried to create objects in fixtures, setup methods, in test function itself, using async_to_sync on WebsocketCommunicator
. However, none of those worked.
If I create a user in a fixture and save it correctly gets an id. However, in my consumer Django does not see that User in the database. And treat it like an anonymous user.
I have a temporary token which I use to authenticate a user on websocket.connect.
@pytest.fixture
def room():
room = generate_room()
room.save()
return room
@pytest.fixture
def room_with_user(room, normal_user):
room.users.add(normal_user)
yield room
room.users.remove(normal_user)
@pytest.fixture
def normal_user():
user = generate_user()
user.save()
return user
@pytest.mark.django_db
class TestConnect:
@pytest.mark.asyncio
async def test_get_connected_client(self, path, room_with_user, temp_token):
assert get_path(room_with_user.id) == path
communicator = QSWebsocketCommunicator(application, path, query_string=get_query_string(temp_token))
connected, subprotocol = await communicator.connect()
assert connected
await communicator.disconnect()
Consumer:
class ChatConsumer(JsonWebsocketConsumer):
def connect(self):
# Called on connection. Either call
self.user = self.scope['user']
self.room_id = self.scope['url_route']['kwargs']['room_id']
group = f'room_{self.room_id}'
users = list(User.objects.all()) # no users here
self.group_name = group
if not (self.user is not None and self.user.is_authenticated):
return self.close({'Error': 'Not authenticated user'})
try:
self.room = Room.objects.get(id=self.room_id, users__id=self.user.id)
except ObjectDoesNotExist:
return self.close({'Error': 'Room does not exists'})
# Send success response
self.accept()
# Save user as active
self.room.active_users.add(self.user)
My authentication Middleware
class OAuthTokenAuthMiddleware:
"""
Custom middleware that takes Authorization header and read OAuth token from it.
"""
def __init__(self, inner):
# Store the ASGI application we were passed
self.inner = inner
def __call__(self, scope):
temp_token = self.get_token(scope)
scope['user'] = self.validate_token(temp_token)
return self.inner(scope)
@staticmethod
def get_token(scope) -> str:
return url_parse.parse_qs(scope['query_string'])[b'token'][0].decode("utf-8")
@staticmethod
def validate_token(token):
try:
token = TemporaryToken.objects.select_related('user').get(token=token)
if token.is_active():
token.delete()
return token.user
else:
return AnonymousUser()
except ObjectDoesNotExist:
return AnonymousUser()
And custom WebsocketCommunicator which accepts query_string in order to include my one time token
class QSWebsocketCommunicator(WebsocketCommunicator):
def __init__(self, application, path, headers=None, subprotocols=None,
query_string: Optional[Union[str, bytes]]=None):
if isinstance(query_string, str):
query_string = str.encode(query_string)
self.scope = {
"type": "websocket",
"path": path,
"headers": headers or [],
"subprotocols": subprotocols or [],
"query_string": query_string or ''
}
ApplicationCommunicator.__init__(self, application, self.scope)
My question is how can I create User, Room, etc. objects in tests/fixtures so that I can access them in Django consumer.
Or do you have another idea how can I overcome this?
It's pretty much impossible to reproduce your issue using the code you've provided. Read about How to create a Minimal, Complete, and Verifiable example. However, I suppose that you should use real transactions in your test as the plain pytest.mark.django_db
will skip the transactions and not store any data in the database per se. A working example:
# routing.py
from django import http
from django.conf.urls import url
from django.contrib.auth.models import User
from channels.routing import ProtocolTypeRouter, URLRouter
from channels.generic.websocket import JsonWebsocketConsumer
class ChatConsumer(JsonWebsocketConsumer):
def connect(self):
self.user = self.scope['user']
print('user in scope, set by middleware:', self.user)
users = list(User.objects.all()) # no users here
print('all users in chat consumer:', users)
if not (self.user is not None and self.user.is_authenticated):
return self.close({'Error': 'Not authenticated user'})
# Send success response
self.accept()
class OAuthTokenAuthMiddleware:
def __init__(self, inner):
# Store the ASGI application we were passed
self.inner = inner
def __call__(self, scope):
token = self.get_token(scope)
print('token in middleware:', token)
scope['user'] = User.objects.get(username=token)
return self.inner(scope)
@staticmethod
def get_token(scope) -> str:
d = http.QueryDict(scope['query_string'])
return d['token']
APP = ProtocolTypeRouter({
'websocket': OAuthTokenAuthMiddleware(URLRouter([url(r'^websocket/$', ChatConsumer)])),
})
Sample fixture that creates a user with username spam
:
@pytest.fixture(scope='function', autouse=True)
def create_user():
with transaction.atomic():
User.objects.all().delete()
user = User.objects.create_user(
'spam', '[email protected]', password='eggs',
first_name='foo', last_name='bar'
)
return user
Now, I mark the test as transactional one, meaning that each query is actually committed. Now the test user is stored into database and the queries made in middleware/consumer can actually return something meaningful:
@pytest.mark.django_db(transaction=True)
@pytest.mark.asyncio
async def test_get_connected_client():
app = OAuthTokenAuthMiddleware(URLRouter([url(r'^websocket/$', ChatConsumer)]))
communicator = QSWebsocketCommunicator(app, '/websocket/', query_string='token=spam')
connected, subprotocol = await communicator.connect()
assert connected
await communicator.disconnect()
Running test test yields the desired result:
$ pytest -vs
================================== test session starts =================================
platform darwin -- Python 3.6.3, pytest-3.4.0, py-1.5.2, pluggy-0.6.0 -- /Users/hoefling/.virtualenvs/stackoverflow/bin/python
cachedir: .pytest_cache
Django settings: spam.settings (from environment variable)
rootdir: /Users/hoefling/projects/private/stackoverflow/so-49136564/spam, inifile: pytest.ini
plugins: celery-4.1.0, forked-0.2, django-3.1.2, cov-2.5.1, asyncio-0.8.0, xdist-1.22.0, mock-1.6.3, hypothesis-3.44.4
collected 1 item
tests/test_middleware.py::test_get_connected_client Creating test database for alias 'default'...
token in middleware: spam
user in scope: spam
all users in chat consumer: [<User: spam>]
PASSEDDestroying test database for alias 'default'...
=============================== 1 passed in 0.38 seconds ================================
Btw you don't need to hack around the WebsocketCommunicator
anymore since it is now able to deal with query strings, see this issue closed.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With