create_test_tables.py

# -*- coding: utf-8 -*
from __future__ import print_function, division, absolute_import
import click

import copy
import os
import subprocess
import functools

import dj_database_url


@click.command()
@click.argument('xdist_slaves', type=int)
def main(xdist_slaves):
    info = functools.partial(click.secho, fg='green')
    migration_hash = subprocess.check_output('django_migration_hash').strip()
    db_config = dj_database_url.config()

    if not db_config.get('NAME'):
        raise click.ClickException('Database must be specified in DATABASE_URL')

    test_db_base_name = 'test_{}'.format(db_config['NAME'])
    template_db_name = '{base}_{hash}'.format(base=test_db_base_name, hash=migration_hash)

    ret = subprocess.call(['createdb', template_db_name])

    if ret == 0:  # createdb was successful -> database created
        info('Database {} created, running migrations'.format(template_db_name))
        env = copy.copy(os.environ)
        env['DATABASE_URL'] = 'postgres:///{}'.format(template_db_name)
        proc = subprocess.Popen(['django-admin', 'migrate'], env=env)
        ret = proc.wait()
        assert ret == 0

    info('Creating test databases...')

    for gateway_id in range(xdist_slaves):
        name = '{name}_gw{gateway_id}'.format(name=test_db_base_name, gateway_id=gateway_id)
        assert subprocess.call(['dropdb', '--if-exists', name]) == 0
        assert subprocess.call(['createdb', '-T', template_db_name, name]) == 0

        info('Created {}'.format(name))