pytest-django experiment

pytest_django.py

"""pytest plugin for testing django projects

This plugin handlers creating and destroying the test environment and
database.  It is only enabled when DJANGO_SETTINGS_MODULES is set in
either the environment, the ini file or a conftest.py file.  In case
of conftest.py files the django modules are only imported at the last
possible moment and it is also possible to test multiple sites
(multiple django settings modules) in one py.test process.

Similar to Django's TestCase, a transaction is started and rolled back
for each test.  Additionally, the settings are copied before each test
and restored at the end of the test, so it is safe to modify settings
within tests.
"""

import os
import sys

import py
import pytest


################ Main plugin ################


def pytest_addoption(parser):
    group = parser.getgroup("general")
    group._addoption('--no-db',
                     action='store_true', dest='no_db', default=False,
                     help='Run tests without a django database')
    parser.addini('DJANGO_SETTINGS_MODULE',
                  'Django settings module to use by pytest-django')


def pytest_namespace():
    return {'load_fixture': load_fixture, 'urls': urls}


def pytest_sessionstart(session):
    """Prepare the global django test environment

    This is only done if either DJANGO_SETTINGS_MODULE is in the
    environment or is set in the ini file.
    """
    djs = os.environ.get('DJANGO_SETTINGS_MODULE')
    if not djs:
        djs = session.config.getini('DJANGO_SETTINGS_MODULE')
    if djs:
        runner = setup_django_env(djs, not session.config.getvalue('no_db'))
        session.config.pytest_django_runner = runner
        session.config.pytest_django_settings = djs


def pytest_sessionfinish(session, exitstatus):
    """Teardown django test environment"""
    runner = getattr(session.config, 'pytest_django_runner', None)
    if runner:
        del session.config.pytest_django_runner
        del session.config.pytest_django_settings
        capture = py.io.StdCapture()
        runner.teardown_test_environment()
        runner.teardown_databases(runner.pytest_django_old_db_config)
        stdout, stderr = capture.reset()
        sys.stderr.write(stderr)   # XXX


def pytest_runtest_setup(item):
    """Test-specific django test environment setup

    This looks up DJANGO_SETTINGS_MODULE in the local conftest and if
    present will setup django's test environment for this settings
    module, if it is not already setup.

    Additionally this will modify the ROOT_URLCONF if the item has the
    ``@pytest.urls()`` decorator applied to it as well as ensure the
    database is in a clean state.
    """
    if not hasattr(item.config, 'pytest_django_runner'):
        return

    if hasattr(item.obj, 'urls'):
        item.pytest_django_urlconf = py.std.django.conf.settings.ROOT_URLCONF
        py.std.django.conf.settings.ROOT_URLCONF = item.obj.urls
        py.std.django.core.urlresolvers.clear_url_caches()

    # Setup db if not disabled, django unittests do this themself
    if (not item.config.option.no_db and
        not is_django_unittest(item) and
        'transaction_test_case' not in item.keywords):
        item.pytest_django_testcase = make_django_testcase(item)
        item.pytest_django_testcase._pre_setup()


def pytest_runtest_teardown(item):
    """Restore the django database and settings"""
    if hasattr(item, 'pytest_django_testcase'):
        item.pytest_django_testcase._post_teardown()
    elif 'transaction_test_case' in item.keywords:
        for db in py.std.django.db.connections:
            # Flush the database adn close database connections.
            # Django does this by default *before* each test instead
            # of after.
            py.std.django.core.management.call_command(
                'flush', verbosity=0, interactive=False, database=db)
        for conn in py.std.django.db.connections.all():
            conn.close()

    if hasattr(item, 'pytest_django_urlconf'):
        py.std.django.conf.settings.ROOT_URLCONF = item.pytest_django_urlconf
        py.std.django.core.urlresolvers.clear_url_caches()


################ Funcargs ################


def pytest_funcarg__client(request):
    """Return a Django test client instance"""
    check_enabled(request.config)
    return py.std.django.test.client.Client()


def pytest_funcarg__admin_client(request):
    """Return a Django test client logged in as an admin user"""
    check_enabled(request.config)
    try:
        py.std.django.contrib.auth.models.User.objects.get(username='admin')
    except py.std.django.contrib.auth.models.User.DoesNotExist:
        user = py.std.django.contrib.auth.models.User.objects.create_user(
            'admin', 'admin@example.com', 'password')
        user.is_staff = True
        user.is_superuser = True
        user.save()
    client = py.std.django.test.client.Client()
    client.login(username='admin', password='password')
    return client


def pytest_funcarg__rf(request):
    """Return a RequestFactory instance"""
    check_enabled(request.config)
    return py.std.django.test.client.RequestFactory()


def pytest_funcarg__settings(request):
    """Return a Django settings object

    Any changes will be restored after the test has been run.
    """
    check_enabled(request.config)
    old_settings = py.std.copy.deepcopy(py.std.django.conf.settings)

    def restore_settings():
        for setting in dir(old_settings):
            if setting == setting.upper():
                setattr(py.std.django.conf.settings,
                        setting, getattr(old_settings, setting))

    request.addfinalizer(restore_settings)
    return py.std.django.conf.settings


def pytest_funcarg__live_server(request):
    """See Django's LiveServerTestCase"""
    check_enabled(request.config)
    if not hasattr(py.std.django.test.testcases, 'LiveServerThread'):
        pytest.fail('live_server not supported in django < 1.4')

    def setup_live_server():
        return LiveServer(*live_server_ports())

    def teardown_live_server(live_server):
        live_server.thread.join()

    return request.cached_setup(setup=setup_live_server,
                                teardown=teardown_live_server,
                                scope='session')


################ Helper functions ################


transaction_test_case = pytest.mark.transaction_test_case


def do_django_imports():
    """Required django imports for the plugin

    Since django is properly messed up with it's global state and code
    execution at import it is almost impossible to import any part of
    django without having DJANGO_SETTINGS_MODULE set.  Since this
    plugin wants to work without this environment variable set we need
    to delay all django imports.

    For this we access django using py.std.django which does a lazy
    import.  The limitation of py.std is that only the top level
    package is imported, but by importing the sub-packages explicitly
    we let the import machinery create all the sub-module references
    and py.std.django will have all required sub-modules available.
    """
    import django.conf
    import django.contrib.auth.models
    import django.core.management
    import django.core.urlresolvers
    import django.db.backends.util
    import django.test
    import django.test.client
    import django.test.simple
    import django.test.testcases


def load_fixture(fixture):
    """Loads a fixture, useful for loading fixtures in funcargs.

    Example:

        def pytest_funcarg__articles(request):
            pytest.load_fixture('test_articles')
            return Article.objects.all()
    """
    if 'DJANGO_SETTINGS_MODULE' not in os.environ:
        pytest.fail('pytest-django not enabled, check DJANGO_SETTINGS_MODULE')
    py.std.django.core.management.call_command('loaddata',
                                               fixture, verbosity=1)


def urls(urlconf):
    """Decorator to change the URLconf for a particular test

    This is similar to the `urls` attribute on Django's `TestCase`.

    Example::

       @pytest.urls('myapp.test_urls')
       def test_something(client):
           assert 'Success!' in client.get('/some_path/')
    """
    if 'DJANGO_SETTINGS_MODULE' not in os.environ:
        pytest.fail('pytest-django not enabled, check DJANGO_SETTINGS_MODULE')

    def wrapper(function):
        function.urls = urlconf
        return function

    return wrapper


def setup_django_env(settings_name, db=True):
    """Setup the global django test environment

    This will disable south, setup the django test environment and
    setup the django database.  If django was already loaded it will
    fail.  It will also respect the --no_db option.

    This does many django imports behind your back.

    Return the django test suite runner
    """
    for name in sys.modules.keys():
        if name.startswith('django.'):
            raise RuntimeError('Django already loaded before env setup')
    os.environ['DJANGO_SETTINGS_MODULE'] = settings_name
    do_django_imports()

    # Disable south migrations
    commands = py.std.django.core.management.get_commands()
    commands['syncdb'] = 'django.core'

    # Setup django test env and db
    runner = py.std.django.test.simple.DjangoTestSuiteRunner(interactive=False)
    if not db:
        def cursor_wrapper(*args, **kwargs):
            pytest.fail('Database access blocked by --no-db')

        def fake_setup():
            py.std.django.db.backends.util.CursorWrapper = cursor_wrapper

        runner.setup_databases = fake_setup
        runner.teardown_databases = lambda x: None
    runner.setup_test_environment()
    runner.pytest_django_old_db_config = runner.setup_databases()

    py.std.django.conf.settings.DEBUG_PROPAGATE_EXCEPTIONS = True

    return runner


def make_django_testcase(item):
    """Retuns a django unittest instance for a test item"""
    if 'transaction_test_case' in item.keywords:
        cls = py.std.django.test.TransactionTestCase
    elif item.config.option.no_db:
        cls = py.std.django.test.TestCase
        cls._fixture_setup = lambda self: None
    else:
        cls = py.std.django.test.TestCase
    return cls(methodName='__init__')


def is_django_unittest(item):
    """Returns True if the item is a Django test case, otherwise False"""
    TestCase = getattr(py.std.django.test,
                       'SimpleTestCase', py.std.django.test.TestCase)
    return (hasattr(item.obj, 'im_class') and
            issubclass(item.obj.im_class, TestCase))


def check_enabled(config):
    if not hasattr(config, 'pytest_django_runner'):
        pytest.fail('pytest-django not enabled, check DJANGO_SETTINGS_MODULE')


class LiveServer(object):
    """Helper class for the live_server funcarg"""

    def __init__(self, host, possible_ports):
        connections_override = {}
        for conn in py.std.django.db.connections.all():
            # SQLite in-memory databases need the connection in the
            # server thread
            if (conn.settings_dict['ENGINE'] == 'django.db.backends.sqlite3'
                and conn.settings_dict['NAME'] == ':memory:'):
                conn.allow_thread_sharing = True
                connections_override[conn.alias] = conn
        self.thread = py.std.django.test.testcases.LiveServerThread(
            host, possible_ports, connections_override)
        self.thread.daemon = True
        self.thread.start()
        self.thread.is_ready.wait()
        if self.thread.error:
            raise self.thread.error

    def __unicode__(self):
        return 'http://%s:%s' % (self.thread.host, self.thread.port)

    def __repr__(self):
        return '<LiveServer listenting at %s>' % unicode(self)

    def __add__(self, other):
        # Support string concatenation
        return unicode(self) + other


def live_server_ports():
    # This code is copy-pasted from django/test/testcases.py
    specified_address = os.environ.get('DJANGO_LIVE_TEST_SERVER_ADDRESS',
                                       'localhost:8081')

    # The specified ports may be of the form '8000-8010,8080,9200-9300'
    # i.e. a comma-separated list of ports or ranges of ports, so we break
    # it down into a detailed list of all possible ports.
    possible_ports = []
    try:
        host, port_ranges = specified_address.split(':')
        for port_range in port_ranges.split(','):
            extremes = map(int, port_range.split('-'))
            assert len(extremes) in [1, 2]
            if len(extremes) == 1:  # Port range of the form '8000'
                possible_ports.append(extremes[0])
            else:                   # Port range of the form '8000-8010'
                for port in range(extremes[0], extremes[1] + 1):
                    possible_ports.append(port)
    except Exception:
        raise Exception('Invalid address ("%s") for live server.' %
                        specified_address)
    return (host, possible_ports)