Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Django channels 2, accessing db in tests

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?

like image 290
Quba Avatar asked Mar 06 '23 18:03

Quba


1 Answers

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.

like image 142
hoefling Avatar answered Mar 11 '23 11:03

hoefling