Run MongoDB inside python for integration testing

mongod.py

from dataclasses import dataclass
import tempfile
import os
import subprocess
import traceback
import logging
from contextlib import contextmanager
import time
from multiprocessing import Manager, Process
import shutil
import json
import argparse

log = logging.getLogger(__file__)

def launch_mongod(shared, port, address, dir) -> None:
    """Execute the mongoDB server in a subprocess."""
    arguments = [
        '--dbpath', dir,
        '--wiredTigerCacheSizeGB', '1',
        '--port', str(port),
        '--bind_ip', address,
        '--pidfilepath', os.path.join(dir, 'pid')
    ]

    kwargs = dict(
        args=' '.join(['mongod'] + arguments),
        stdout=subprocess.PIPE,
        bufsize=1,
        stderr=subprocess.STDOUT,
    )

    shared['db_pid_path'] = os.path.join(dir, 'pid')

    with subprocess.Popen(**kwargs, shell=True) as proc:
        try:
            shared['pid'] = proc.pid

            if 'dp_pid' not in shared:
                try:
                    shared['dp_pid'] = int(open(os.path.join(dir, 'pid'), 'r').read())
                    shared['running'] = True
                except FileNotFoundError:
                    pass

            while shared['running']:
                if proc.poll() is None:
                    line = proc.stdout.readline().decode('utf-8')[:-1]
                    log.debug(line)
                else:
                    shared['running'] = False
                    shared['exit'] = proc.returncode
                    log.debug('Stopping mongod popen')

        except Exception:
            log.error(traceback.format_exc())
            shared['running'] = False
            return

@contextmanager
def process_mongod(port, address) -> None:
    """Launch a mongoDB server in parallel. The server is stop on exit"""

    with tempfile.TemporaryDirectory() as dir:
        with Manager() as manager:
            shared = manager.dict()
            shared['running'] = False

            proc = Process(target=launch_mongod, args=(shared, port, address, dir))
            proc.start()

            acc = 0
            while not shared['running'] and acc < 2:
                acc += 0.01
                time.sleep(0.01)

            log.debug('Mongod ready after %f', acc)

            yield proc

            log.debug("Should stop")

            shared['running'] = False
            stop_mongodb(dir)

            log.debug('Stopping mongod process')
            shutil.rmtree(dir)
            proc.kill()


def start_mongodb(port, address, dir):
    """Start the mongoDB server."""

    os.environ['MONGO_PATH'] = dir
    os.environ['MONGO_PORT'] = str(port)
    os.environ['MONGO_ADDRESS'] = str(address)
    mongodb_run_script = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', 'scripts', 'setup.sh'))

    log.debug("Started mongo")
    subprocess.call(args=['bash', mongodb_run_script])


def stop_mongodb(dir):
    """Stop the mongoDB server."""
    subprocess.run(['mongod', '--shutdown', '--dbpath', dir])


@dataclass
class MongoDBConfig:
    port: int
    address: str
    dir: str
    admin: str = 'god'
    admin_password: str = 'god123'


@contextmanager
def mongod(port, address) -> None:
    """Launch a mongoDB server in parallel. The server is stop on exit"""

    with tempfile.TemporaryDirectory() as dir:
        start_mongodb(port, address, dir);

        yield MongoDBConfig(port, address, dir)

        stop_mongodb(dir)


def _add_mongo_user(mongo, user, password, db='amin', role='admin'):
    """Add a user to the mongoDB server."""

    create_user = dict(
        user=user,
        pwd=password,
        roles=[
            dict(role=role, db=db),
        ],
    )

    subprocess.run(['mongo',
        f'{mongo.address}:{str(mongo.port)}',
        '--authenticationDatabase', 'admin',
        '-u', mongo.admin,
        '-p', mongo.admin_password,
        '--eval', f'db.createUser({json.dumps(create_user)})',
    ])


def add_admin_user(mongo, user, password):
    """Add a admin user to the mongoDB server."""
    _add_mongo_user(mongo, user, password, db='admin', role='userAdminAnyDatabase')


def add_orion_user(mongo, user, password):
    """Add a orion user to the mongoDB server."""
    _add_mongo_user(mongo, user, password, db='orion', role='readWrite')


def main():
    """Entry point for tox to setup and shutdown testing databases"""

    parser = argparse.ArgumentParser(description='Setup and shutdown testing databases')
    parser.add_argument('command', type=str, choices=['start', 'stop'], help='Command to execute')
    parser.add_argument('--port', type=int, default=27017, help='Port for the mongoDB server')
    parser.add_argument('--address', type=str, default='localhost', help='Address for the mongoDB server')
    parser.add_argument('--dir', type=str, default='/tmp/orion', help='Directory for the mongoDB server')
    args = parser.parse_args()

    if args.command == 'start':
        start_mongodb(args.port, args.address, args.dir)
    else:
        stop_mongodb(args.dir)

setup.sh

#!/bin/bash

# Setup a mongodb server for testing
# this setup is done my IT when the DB is deployed
# set -evx

PORT=${MONGO_PORT:-"8123"}
ADDRESS=${MONGO_ADDRESS:-"localhost"}
ADMIN=${MONGO_ADMIN:-"god"}
PASSWORD=${MONGO_PASS:-"god123"}
DB_PATH=${MONGO_PATH:-"/tmp/db"}

function start_init_mongod {
    # Start mongodb without Access Control, this is used to insert the admin user
    # the db is then stopped and started with Access Control

    if test -f "$DB_PATH/pid"; then
        pid=$(cat $DB_PATH/pid)
        kill -3 $pid || true
    fi

    rm -rf $DB_PATH
    mkdir -p $DB_PATH

    mongod --dbpath $DB_PATH/ --wiredTigerCacheSizeGB 1 --port $PORT --bind_ip localhost --pidfilepath $DB_PATH/pid &
    sleep 1
}

function start_mongod {
    # start mongodb with Access Control
    mongod --auth --dbpath $DB_PATH/ --wiredTigerCacheSizeGB 1 --port $PORT --bind_ip $ADDRESS --pidfilepath $DB_PATH/pid --fork --logpath $DB_PATH/log.txt
    sleep 1
}

function stop_mongod  {
    echo "$(pwd)"
    pid=$(cat $DB_PATH/pid)
    kill -s TERM $pid
    rm -rf $DB_PATH/pid
    sleep 1
}

function add_admin_user {
    # create an admin user
    # userAdminAnyDatabase: create users, grant & revoke roles, create and modify roles

    CMD=$(cat <<EOM
    use admin
    db.createUser({
        user: "$ADMIN",
        pwd: "$PASSWORD",
        roles: [
            { role: "userAdminAnyDatabase", db: "admin" },
            { role: "readWriteAnyDatabase", db: "admin" }
        ]
    })
EOM
    )

    echo "$CMD" | mongo --port $PORT
}

function add_user {
    # Create a user for the orion database

    username=$1
    password=$2

    CMD=$(cat << EOM
    use orion
    db.createUser({
        user: "$username",
        pwd: "$password",
        roles: [
            { role: "readWrite", db: "orion" }
        ]
    })
EOM
    )

    echo "$CMD" | mongo "mongodb://$ADDRESS:$PORT" --authenticationDatabase "admin" -u $ADMIN -p $PASSWORD
}


function ensure_indexes {
    # User will have limited access to the collection
    # so orion's client cannot do this

    CMD=$(cat <<EOM
    use orion

EOM
    )

    echo "$CMD" | mongo --port $PORT
}

function launch {
    # Setup a mongodb for testing
    start_init_mongod
    add_admin_user
    ensure_indexes
    stop_mongod

    # Start mongodb with access control
    start_mongod

    add_user User1 Pass1
    add_user User2 Pass2
    add_user User3 Pass3
}

export MONGO_RUNNING="${DB_PATH}"
launch