init_db.py

# Copyright 2020 QuantStack
# Distributed under the terms of the Modified BSD License.

import random
import uuid

from quetz.config import Config
from quetz.database import get_session, init_db
from quetz.db_models import (
    ApiKey,
    Channel,
    ChannelMember,
    Identity,
    Package,
    PackageMember,
    Profile,
    User,
)


def init_test_db():
    config = Config()
    init_db(config.sqlalchemy_database_url)
    db = get_session(config.sqlalchemy_database_url)

    testUsers = []

    try:
        for index, username in enumerate(['alice', 'bob', 'carol', 'dave']):
            user = User(id=uuid.uuid4().bytes, username=username)

            identity = Identity(
                provider='dummy',
                identity_id=str(index),
                username=username,
            )

            profile = Profile(name=username.capitalize(), avatar_url='/avatar.jpg')

            user.identities.append(identity)
            user.profile = profile
            db.add(user)
            testUsers.append(user)

        for channel_index in range(30):
            channel = Channel(
                name=f'channel{channel_index}',
                description=f'Description of channel{channel_index}',
                private=False,
            )

            for package_index in range(random.randint(5, 100)):
                package = Package(
                    name=f'package{package_index}',
                    description=f'Description of package{package_index}',
                )
                channel.packages.append(package)

                test_user = testUsers[random.randint(0, len(testUsers) - 1)]
                package_member = PackageMember(
                    package=package, channel=channel, user=test_user, role='owner'
                )

                db.add(package_member)

            if channel_index == 0:
                package = Package(name='xtensor', description='Description of xtensor')
                channel.packages.append(package)

                test_user = testUsers[random.randint(0, len(testUsers) - 1)]
                package_member = PackageMember(
                    package=package, channel=channel, user=test_user, role='owner'
                )

                db.add(package_member)

                # create API key
                key = 'E_KaBFstCKI9hTdPM7DQq56GglRHf2HW7tQtq6si370'

                key_user = User(id=uuid.uuid4().bytes)

                api_key = ApiKey(
                    key=key, description='test API key', user=key_user, owner=test_user
                )
                db.add(api_key)

                key_package_member = PackageMember(
                    user=key_user,
                    channel_name=channel.name,
                    package_name=package.name,
                    role='maintainer',
                )
                db.add(key_package_member)

            db.add(channel)

            channel_member = ChannelMember(
                channel=channel,
                user=testUsers[random.randint(0, len(testUsers) - 1)],
                role='owner',
            )

            db.add(channel_member)
        db.commit()
    finally:
        db.close()


if __name__ == '__main__':
    init_test_db()