utils.py

import os
import unittest
from urllib.parse import urljoin

import requests

_DEFAULT_BASE_URL = 'http://hubapp'
_DEFAULT_ADMIN_USERNAME = 'admin'
_DEFAULT_ADMIN_PASSWORD = 'password'


class BaseUrlSession(requests.Session):
    def __init__(self, base_url=None, *args, **kwargs):
        self.base_url = base_url
        super().__init__(*args, **kwargs)

    def request(self, method, url, *args, **kwargs):
        url = urljoin(self.base_url, url)
        return super().request(method, url, *args, **kwargs)


class HubTestCase(unittest.TestCase):
    require_api = False
    require_backoffice = False
    require_frontend = False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # get base URL from environment variable
        self.base_url = os.getenv('BASE_URL', _DEFAULT_BASE_URL)

        # ensure base_url ends with a slash
        if not self.base_url.endswith('/'):
            self.base_url += '/'

        self.session = BaseUrlSession(base_url=self.base_url)

    def setUp(self, *args, **kwargs):
        super().setUp(*args, **kwargs)
        if self.require_api and os.getenv('SERVE_API') not in ['y', 'Y', 'yes']:
            self.skipTest('environment: SERVE_API=no')
        if self.require_backoffice and os.getenv('SERVE_BACKOFFICE') not in ['y', 'Y', 'yes']:
            self.skipTest('environment: SERVE_BACKOFFICE=no')
        if self.require_frontend and os.getenv('SERVE_FRONTEND') not in ['y', 'Y', 'yes']:
            self.skipTest('environment: SERVE_FRONTEND=no')

    def tearDown(self, *args, **kwargs):
        super().tearDown(*args, **kwargs)

        # close the session to avoid warning
        self.session.close()


class HubApiTestCase(HubTestCase):
    require_api = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._admin_session = None

    def tearDown(self, *args, **kwargs):
        super().tearDown(*args, **kwargs)

        # close an open admin session to avoid warning
        if self._admin_session is not None:
            self._admin_session.close()

    @property
    def admin_session(self):
        if self._admin_session is not None:
            return self._admin_session

        try:
            session = BaseUrlSession(base_url=self.base_url)
            login_success = self.api_login(
                session,
                username=os.getenv('DJANGO_CREATE_ADMIN_USERNAME', _DEFAULT_ADMIN_USERNAME),
                password=os.getenv('DJANGO_CREATE_ADMIN_PASSWORD', _DEFAULT_ADMIN_PASSWORD),
            )

            if login_success:
                self._admin_session = session
                return session
            else:
                self.fail('Cannot login into admin session via API.')

        except Exception:
            session.close()
            raise

    def api_login(self, session: requests.Session, username: str, password: str) -> bool:
        response = session.post('api/auth/get-token', data={'username': username, 'password': password})
        if not response.ok:
            self.fail(f'API Auth Get-Token returned error {response.status_code}: {response.text}')

        # extract token
        token = response.json().get('token')
        if token is None:
            return False

        # store token
        session.token = token
        session.headers['Authorization'] = 'Token ' + token

        return True

    def logout(self, session: requests.Session):
        self.headers['Authorization'] = None