tests.py

# Core
import os
import sys
import shutil
import tempfile
import functools
import contextlib
from glob import glob
from datetime import datetime
from importlib import import_module

# 3rd Party
import nose
import traceback
import mongoengine
from click import echo
from pymongo import MongoClient
from nose.tools import with_setup
from click.testing import CliRunner

# backward / forward support for reload
from six.moves import reload_module
from nose.plugins.skip import SkipTest


# Local
from monarch import cli, create_package_if_necessary

mongo_port = int(os.environ.get('MONARCH_MONGO_DB_PORT', 27017))

TEST_ENVIRONEMNTS = {
    'test': {
        'db_name': 'test_monarch',
        'host': 'localhost:{}'.format(mongo_port)
    },
    'from_test': {
        'db_name': 'from_monarch_test',
        'host': 'localhost:{}'.format(mongo_port)
    },
    'to_test': {
        'db_name': 'to_monarch_test',
        'host': 'localhost:{}'.format(mongo_port)
    },
}


def generate_mongo_uri(environ):
    """this assumes a environ with db_name and host is well formed in the test_environment"""
    return "mongodb:://{}/{}".format(environ['host'], environ['db_name'])

BACKUPS = {
    'LOCAL': {
        'backup_dir': 'path_to_backups',
    },
}


TEST_CONFIG = """
# monarch settings file, generated by monarch init
# feal free to edit it with your application specific settings

ENVIRONMENTS = {}

BACKUPS = {}

""".format(TEST_ENVIRONEMNTS, BACKUPS)

TEST_MIGRATION = """
from monarch import MongoBackedMigration

class {migration_class_name}(MongoBackedMigration):

    def run(self):
        print("running a migration with no failure")
"""

TEST_FAILED_MIGRATION = """
from monarch import MongoBackedMigration

class {migration_class_name}(MongoBackedMigration):

    def run(self):
        print("running a migration with failure")
        raise Exception('Yikes we messed up the database -- oh nooooooooo')
"""

def assert_normal_execution(result):
    if result.exit_code != 0:
        echo('exit_code: {}'.format(result.exit_code))
        echo('output: {}'.format(result.output))
        echo('exception: {}'.format(result.exception))
        _, _, tb = result.exc_info
        echo('trace: {}'.format(traceback.print_tb(tb)))

        assert False
    else:
        assert True


def eq_(this, that):
    assert this == that, "{} did not equal {}".format(this, that)


def no_op():
    pass


@contextlib.contextmanager
def isolated_filesystem_with_path():
    """A context manager that creates a temporary folder and changes
    the current working directory to it for isolated filesystem tests.

    The modification here is that it adds itself to the path

    """
    cwd = os.getcwd()
    t = tempfile.mkdtemp()
    os.chdir(t)
    sys.path.insert(1, t)
    try:
        yield t
    finally:
        os.chdir(cwd)
        sys.path.remove(t)
        try:
            shutil.rmtree(t)
        except (OSError, IOError):
            pass


def requires_mongoengine(func):
    @functools.wraps(func)
    def wrapper(*args, **kw):
        if mongoengine is None:
            raise SkipTest("mongoengine is not installed")
        return func(*args, **kw)

    return wrapper


def clear_mongo_databases():
    for env_name in TEST_ENVIRONEMNTS:
        env = TEST_ENVIRONEMNTS[env_name]
        host, port = env['host'].split(':')
        client = MongoClient(host=host, port=int(port))
        client.drop_database(env['db_name'])


def initialize_monarch(working_dir, backup_dir=None):
    migration_path = os.path.join(os.path.abspath(working_dir), 'migrations')
    create_package_if_necessary(migration_path)
    settings_file = os.path.join(os.path.abspath(working_dir), 'migrations/settings.py')
    with open(settings_file, 'w') as f:
        if backup_dir:
            f.write(TEST_CONFIG.replace('path_to_backups', backup_dir))
        else:
            f.write(TEST_CONFIG)

    m = import_module('migrations')
    reload_module(m)
    s = import_module('migrations.settings')
    reload_module(s)


def ensure_current_migrations_module_is_loaded():
    # everytime within the same python process we add migrations we need to reload the migrations module
    # for it could be cached from a previous test
    if sys.version_info >= (3, 3):
        from importlib import invalidate_caches
        invalidate_caches()
    m = import_module('migrations')
    reload_module(m)


def test_create_migration():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        initialize_monarch(working_dir)
        result = runner.invoke(cli, ['generate', 'add_indexes'])
        new_files_generated = glob(working_dir + '/*/*migration.py')
        assert len(new_files_generated) == 1
        file_name = new_files_generated[0]
        assert 'add_indexes_migration.py' in file_name
        assert os.path.getsize(file_name) > 0
        assert result.exit_code == 0


def test_initialization():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        result = runner.invoke(cli, ['init'])
        settings_file = os.path.join(working_dir, 'migrations/settings.py')
        assert os.path.getsize(settings_file) > 0
        assert result.exit_code == 0


def first_migration(working_dir):
    new_files_generated = glob(working_dir + '/*/*migration.py')
    assert len(new_files_generated) == 1
    return new_files_generated[0]


def establish_connection(env_name):
    s = import_module('migrations.settings')
    env = s.ENVIRONMENTS[env_name]
    mongoengine.connect(env['db_name'], host=generate_mongo_uri(env))


def get_db(env):
    host, port = env['host'].split(':')
    client = MongoClient(host=host, port=int(port))
    return client[env['db_name']]


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_run_migration():
    runner = CliRunner()
    with isolated_filesystem_with_path() as cwd:
        initialize_monarch(cwd)
        runner.invoke(cli, ['generate', 'add_column_to_user_table'])
        # Update Migration Template with a *proper* migration
        current_migration = first_migration(cwd)
        class_name = "{}Migration".format('AddColumnToUserTable')
        with open(current_migration, 'w') as f:
            f.write(TEST_MIGRATION.format(migration_class_name=class_name))

        ensure_current_migrations_module_is_loaded()

        result = runner.invoke(cli, ['migrate', 'test'])
        # echo('output: {}'.format(result.output))
        # echo('exception: {}'.format(result.exception))
        assert result.exit_code == 0


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_failed_migration():
    runner = CliRunner()
    with isolated_filesystem_with_path() as cwd:
        initialize_monarch(cwd)
        result = runner.invoke(cli, ['generate', 'add_account_table'])
        # echo('output: {}'.format(result.output))
        # echo('exception: {}'.format(result.exception))

        # Update Migration Template with a *proper* migration
        current_migration = first_migration(cwd)
        class_name = "{}Migration".format('AddAccountTable')
        with open(current_migration, 'w') as f:
            f.write(TEST_FAILED_MIGRATION.format(migration_class_name=class_name))

        ensure_current_migrations_module_is_loaded()

        result = runner.invoke(cli, ['migrate', 'test'])
        # echo('output: {}'.format(result.output))
        # echo('exception: {}'.format(result.exception))
        assert result.exit_code == -1


def populate_database(env_name):
    from_db = get_db(TEST_ENVIRONEMNTS[env_name])
    from_fishes = from_db.fishes
    fish = {'name': "Red Fish"}
    from_fishes.insert(fish)
    assert from_fishes.count() == 1


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_copy_db():
    runner = CliRunner()
    with isolated_filesystem_with_path() as cwd:
        initialize_monarch(cwd)

        populate_database('from_test')

        to_db = get_db(TEST_ENVIRONEMNTS['to_test'])
        to_fishes = to_db.fishes

        assert to_fishes.count() == 0

        result = runner.invoke(cli, ['copy_db', 'from_test:to_test'], input="y\ny\n")
        echo('trd_a output: {}'.format(result.output))
        echo('trd_a exception: {}'.format(result.exception))

        assert to_fishes.count() == 1


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_list_migrations():
    runner = CliRunner()

    with isolated_filesystem_with_path() as working_dir:
        initialize_monarch(working_dir)
        for migration_name in ['add_indexes', 'add_user_table', 'add_account_table']:
            runner.invoke(cli, ['generate', migration_name])

        ensure_current_migrations_module_is_loaded()

        result = runner.invoke(cli, ['list_migrations', 'test'])

        assert result.exit_code == 0


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_one_off_migration():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        initialize_monarch(working_dir)

        runner.invoke(cli, ['generate', 'add_column_to_user_table'])
        # Update Migration Template with a *proper* migration
        current_migration = first_migration(working_dir)
        class_name = "{}Migration".format('AddColumnToUserTable')
        with open(current_migration, 'w') as f:
            f.write(TEST_MIGRATION.format(migration_class_name=class_name))

        for migration_name in ['add_indexes', 'add_user_table', 'add_account_table']:
            runner.invoke(cli, ['generate', migration_name])

        ensure_current_migrations_module_is_loaded()

        list_result = runner.invoke(cli, ['list_migrations', 'test'])

        a_migration_to_run = list_result.output.split('\n')[4].split(' ')[0]

        result = runner.invoke(cli, ['migrate_one', a_migration_to_run, 'test'])
        assert result.exit_code == 0

        # do it again -- you can with migrate_one
        result = runner.invoke(cli, ['migrate_one', a_migration_to_run, 'test'])
        assert result.exit_code == 0


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_backup_database():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        backup_dir = os.path.join(working_dir, 'backups')
        os.mkdir(backup_dir)

        initialize_monarch(working_dir, backup_dir=backup_dir)
        populate_database('from_test')

        result = runner.invoke(cli, ['backup', 'from_test'])

        assert result.exit_code == 0
        assert len([name for name in os.listdir(backup_dir)]) == 1


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_list_backups():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        backup_dir = os.path.join(working_dir, 'backups')
        os.mkdir(backup_dir)

        initialize_monarch(working_dir, backup_dir=backup_dir)
        populate_database('from_test')

        runner.invoke(cli, ['backup', 'from_test'])
        runner.invoke(cli, ['backup', 'from_test'])

        result = runner.invoke(cli, ['list_backups'])

        echo('tlb output: {}'.format(result.output))
        echo('tlb exception: {}'.format(result.exception))

        assert result.exit_code == 0
        assert "from_monarch_test__{}.dmp.zip".format(datetime.utcnow().strftime("%Y_%m_%d")) in result.output
        assert "from_monarch_test__{}_2.dmp.zip".format(datetime.utcnow().strftime("%Y_%m_%d")) in result.output


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_restore_database():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        backup_dir = os.path.join(working_dir, 'backups')
        os.mkdir(backup_dir)

        initialize_monarch(working_dir, backup_dir=backup_dir)
        populate_database('from_test')

        result = runner.invoke(cli, ['backup', 'from_test'])
        assert result.exit_code == 0

        migration_name = "from_monarch_test__{}.dmp.zip".format(datetime.utcnow().strftime("%Y_%m_%d"))
        result = runner.invoke(cli, ['restore', "{}:to_test".format(migration_name)], input="y\ny\n")
        assert result.exit_code == 0

        # Check to see if the database that was imported has the right data
        to_db = get_db(TEST_ENVIRONEMNTS['to_test'])
        to_fishes = to_db.fishes
        assert to_fishes.count() == 1


def test_create_query_set():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        initialize_monarch(working_dir)
        result = runner.invoke(cli, ['generate_query_set', 'last_three_weeks'])
        new_files_generated = glob(working_dir + '/querysets/*_queryset.py')
        assert len(new_files_generated) == 1
        file_name = new_files_generated[0]
        assert 'last_three_weeks_queryset.py' in file_name
        assert os.path.getsize(file_name) > 0
        assert result.exit_code == 0


def set_up_from_db_for_queryset_tests():
    from_db = get_db(TEST_ENVIRONEMNTS['from_test'])

    # Dogs
    from_dogs = from_db.dogs
    rex = {'name': "Rex", 'type': "Awesome"}
    rover = {'name': "Rover", 'type': "Silly"}
    rex_id = from_dogs.insert(rex)
    rover_id = from_dogs.insert(rover)

    eq_(from_dogs.count(), 2)

    # Dog Houses
    from_dog_houses = from_db.dog_houses
    rex_house = {'name': 'Rex House', 'dog_id': rex_id}
    rover_house = {'name': 'Rover House', 'dog_id': rover_id}
    from_dog_houses.insert(rex_house)
    from_dog_houses.insert(rover_house)

    eq_(from_dog_houses.count(), 2)

    # Cats
    from_cats = from_db.cats
    muffy = {'name': 'Muffy'}
    puffy = {'name': 'Puffy'}
    from_cats.insert(muffy)
    from_cats.insert(puffy)

    eq_(from_cats.count(), 2)


def generate_and_import_queryset_file(cwd, runner, template, file_name):
    # Create the query_set file
    result = runner.invoke(cli, ['generate_query_set', file_name])
    assert_normal_execution(result)
    #current query_set
    new_files_generated = glob(cwd + '/querysets/*queryset.py')
    assert len(new_files_generated) == 1
    current_query_set = new_files_generated[0]
    with open(current_query_set, 'w') as f:
        f.write(template)

    # ensure that queryset is loaded
    m = import_module('querysets')
    reload_module(m)


V1_TEST_QUERY_SET = """
from monarch import QuerySet
from click import echo

class AwesomeDogsQuerySet(QuerySet):

    def run(self):

        awesome_dogs = self.database.dogs.find({"type": "Awesome"})
        awesome_dog_ids = [dog['_id'] for dog in awesome_dogs]
        echo("awesome dog ids: {}".format(awesome_dog_ids))

        self.dump_collection('dogs', {"_id": {"$in": awesome_dog_ids}})
        self.dump_collection('dog_houses', {"dog_id": {"$in": awesome_dog_ids}})

"""

@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_basic_query_set_with_copydb():
    runner = CliRunner()
    with isolated_filesystem_with_path() as cwd:
        initialize_monarch(cwd)

        set_up_from_db_for_queryset_tests()

        to_db = get_db(TEST_ENVIRONEMNTS['to_test'])
        to_dogs = to_db.dogs
        to_dog_houses = to_db.dog_houses

        assert to_dogs.count() == 0
        assert to_dog_houses.count() == 0
        assert to_db.cats.count() == 0

        generate_and_import_queryset_file(cwd, runner, V1_TEST_QUERY_SET, 'awesome_dogs')

        result = runner.invoke(cli, ['copy_db', 'from_test:to_test', '--query-set=AwesomeDogsQuerySet'], input="y\ny\n")

        assert_normal_execution(result)

        eq_(to_dogs.count(), 1)
        eq_(to_dog_houses.count(), 1)

        # If you do not want to include cats, you need to either setup the include or exclude properties on
        # QuerySet
        eq_(to_db.cats.count(), 2)

@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_basic_query_set_with_backup():
    runner = CliRunner()
    with isolated_filesystem_with_path() as cwd:
        backup_dir = os.path.join(cwd, 'backups')
        os.mkdir(backup_dir)

        initialize_monarch(cwd, backup_dir=backup_dir)

        set_up_from_db_for_queryset_tests()

        generate_and_import_queryset_file(cwd, runner, V1_TEST_QUERY_SET, 'awesome_dogs')

        result = runner.invoke(cli, ['backup', 'from_test', '--query-set=AwesomeDogsQuerySet'])

        assert_normal_execution(result)


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_list_query_sets():
    runner = CliRunner()
    with isolated_filesystem_with_path() as working_dir:
        backup_dir = os.path.join(working_dir, 'backups')
        os.mkdir(backup_dir)

        initialize_monarch(working_dir, backup_dir=backup_dir)

        set_up_from_db_for_queryset_tests()

        generate_and_import_queryset_file(working_dir, runner, V1_TEST_QUERY_SET, 'awesome_dogs')

        result = runner.invoke(cli, ['list_query_sets'])

        echo('tlb output: {}'.format(result.output))
        echo('tlb exception: {}'.format(result.exception))

        assert result.exit_code == 0
        assert "AwesomeDogsQuerySet" in result.output


PROMPT_QUERY_SET = """
from monarch import QuerySet
from click import echo, prompt

class AwesomeDogsQuerySet(QuerySet):

    def run(self):

        dog_name = prompt('Who is your favorite dog?')

        awesome_dogs = self.database.dogs.find({"name": dog_name})
        awesome_dog_ids = [dog['_id'] for dog in awesome_dogs]
        echo("awesome dog ids: {}".format(awesome_dog_ids))

        self.dump_collection('dogs', {"_id": {"$in": awesome_dog_ids}})
        self.dump_collection('dog_houses', {"dog_id": {"$in": awesome_dog_ids}})

"""

@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_prompt_query_set():
    runner = CliRunner()
    with isolated_filesystem_with_path() as cwd:
        initialize_monarch(cwd)

        set_up_from_db_for_queryset_tests()

        to_db = get_db(TEST_ENVIRONEMNTS['to_test'])
        to_dogs = to_db.dogs
        to_dog_houses = to_db.dog_houses

        assert to_dogs.count() == 0
        assert to_dog_houses.count() == 0

        generate_and_import_queryset_file(cwd, runner, PROMPT_QUERY_SET, 'awesome_dogs')

        result = runner.invoke(cli, ['copy_db', 'from_test:to_test', '--query-set=AwesomeDogsQuerySet'], input="y\nRex\ny\n")

        assert_normal_execution(result)

        eq_(to_dogs.count(), 1)
        eq_(to_dog_houses.count(), 1)


EXCLUDE_QUERY_SET = """
from monarch import QuerySet
from click import echo, prompt

class ExcludeCatsQuerySet(QuerySet):

    def exclude(self):
        return ['cats']

    def run(self):

        dog_name = prompt('Who is your favorite dog?')

        awesome_dogs = self.database.dogs.find({"name": dog_name})
        awesome_dog_ids = [dog['_id'] for dog in awesome_dogs]
        echo("awesome dog ids: {}".format(awesome_dog_ids))

        self.dump_collection('dogs', {"_id": {"$in": awesome_dog_ids}})
        self.dump_collection('dog_houses', {"dog_id": {"$in": awesome_dog_ids}})

"""


@requires_mongoengine
@with_setup(clear_mongo_databases, clear_mongo_databases)
def test_query_set_exclude():
    runner = CliRunner()
    with isolated_filesystem_with_path() as cwd:
        initialize_monarch(cwd)

        set_up_from_db_for_queryset_tests()

        to_db = get_db(TEST_ENVIRONEMNTS['to_test'])

        # start with a clean environment
        assert to_db.dogs.count() == 0
        assert to_db.dog_houses.count() == 0
        assert to_db.cats.count() == 0

        generate_and_import_queryset_file(cwd, runner, EXCLUDE_QUERY_SET, 'excludes_cats')

        result = runner.invoke(cli, ['copy_db', 'from_test:to_test', '--query-set=ExcludeCatsQuerySet'], input="y\nRex\ny\n")

        assert_normal_execution(result)

        eq_(to_db.dogs.count(), 1)
        eq_(to_db.dog_houses.count(), 1)
        eq_(to_db.cats.count(), 0)

if __name__ == "__main__":
    nose.run()