testing.py

"""
Test overrides to support Safe Cookies with Test Clients.
"""

from django.test.client import Client


def safe_cookie_test_session_patch():
    """
    Override the Test Client's methods in order to support Safe Cookies.
    If there's a better way to patch this, we should do so.
    """
    if getattr(safe_cookie_test_session_patch, 'has_run', False):
        return

    def using_safe_cookie_data(settings):
        """
        Returns whether or not Safe Cookies is actually being
        used, by checking the middleware settings.
        """
        return (
            'openedx.core.djangoapps.safe_sessions.middleware.SafeSessionMiddleware' in settings.MIDDLEWARE_CLASSES
        )

    ## session_id --> safe_cookie_data ##

    # Override Client.login method to update cookies with safe
    # cookies.
    patched_client_login = Client.login

    def login_with_safe_session(self, **credentials):
        """
        Call the original Client.login method, but update the
        session cookie with a freshly computed safe_cookie_data
        before returning.
        """
        from django.conf import settings
        from django.contrib.auth import SESSION_KEY
        from .middleware import SafeSessionMiddleware

        if not patched_client_login(self, **credentials):
            return False
        if using_safe_cookie_data(settings):
            SafeSessionMiddleware.update_with_safe_session_cookie(self.cookies, self.session[SESSION_KEY])
        return True
    Client.login = login_with_safe_session

    ## safe_cookie_data --> session_id ##

    # Override Client.session so any safe cookies are parsed before
    # use.
    def get_safe_session(self):
        """
        Here, we are duplicating the original Client._session code
        in order to allow conversion of the safe_cookie_data back
        to the raw session_id, if needed.  Since test code may
        access the session_id before it's actually converted,
        we use a try-except clause here to check both cases.
        """
        from django.apps import apps
        from django.conf import settings
        from importlib import import_module
        from .middleware import SafeCookieData, SafeCookieError, SafeSessionMiddleware

        if apps.is_installed('django.contrib.sessions'):
            engine = import_module(settings.SESSION_ENGINE)
            cookie = self.cookies.get(settings.SESSION_COOKIE_NAME, None)
            if cookie:
                session_id = cookie.value
                if using_safe_cookie_data(settings):
                    try:
                        session_id = SafeCookieData.parse(session_id).session_id
                    except SafeCookieError:
                        pass  # The safe cookie hasn't yet been created.
                return engine.SessionStore(session_id)
            else:
                session = engine.SessionStore()
                session.save()
                self.cookies[settings.SESSION_COOKIE_NAME] = session.session_key
                SafeSessionMiddleware.update_with_safe_session_cookie(self.cookies, user_id=None)
                return session
        return {}
    Client.session = property(get_safe_session)

    safe_cookie_test_session_patch.has_run = True