Pulp

import logging

import logging.handlers

import random

import uuid

from pulp.server import initialization

from pulp.server.db import connection

from pulp.server.controllers import consumer as consumer_controller

from pulp.server.managers import factory

_logger = logging.getLogger(__name__)

_logger.setLevel(logging.INFO)

_logger.addHandler(logging.handlers.SysLogHandler(address='/dev/log'))

initialization.initialize()

db = connection.get_database()

consumers_c = db['consumers']

bindings_c = db['consumer_bindings']

profiles_c = db['consumer_unit_profiles']

repos_c = db['repos']

rpa_c = db['repo_profile_applicability']

rcu_c = db['repo_content_units']

rpm_c = db['units_rpm']

erratum_c = db['units_erratum']

distributors_c = db['repo_distributors']

def drop_consumer_collections():

    consumers_c.drop()

    bindings_c.drop()

    profiles_c.drop()

    rpa_c.drop()

def generate_and_set_consumers(num=500, profile_size=1000, applicability_rate=0):

    """

    :param num: number of consumer profiles to generate

    :type num: int

    :param profile_size: number of packages in consumer profile

    :type profile_size: int

    :param applicability_rate: N of applicable packages/N of packages in consumer profile

    :type applicability_rate: float

    """

    consumers = []

    for _ in range(num):

        consumer_id = str(uuid.uuid4())

        _logger.info('Registering consumer %s' % consumer_id)

        register_consumer(consumer_id)

        consumers.append(consumer_id)

    generate_consumer_profile(consumers, profile_size, applicability_rate)

def register_consumer(consumer_id):

    """

    :param consumer_id: Id of a consumer to register with

    :type consumer_id: str

    """

    manager = factory.consumer_manager()

    consumer, certificate = manager.register(consumer_id)

def set_consumer_profile(consumer_id, profile, content_type='rpm'):

    """

    :param consumer_id: Id of a consumer to set a profile for

    :type consumer_id: str

    :param profile:  A profile to associate to consumer

    :type profile: list

    """

    manager = factory.consumer_profile_manager()

    new_profile = manager.create(consumer_id, content_type, profile)

    _logger.info('Profile set')

def bind_consumer(consumer_id, repo_id, distributor_id):

    """

    :param consumer_id: Id of a consumer

    :type consumer_id: str

    :param repo_id: Id of repo to bind to

    :type repo_id: str

    :param distributor_id: Id of repo distributor to bind to

    :type distributor_id: str

    """

    notify_agent = False

    consumer_controller.bind(consumer_id, repo_id, distributor_id, notify_agent, {}, {})

    _logger.info('Consumer bound to %s' % repo_id)

def generate_consumer_profile(consumers, size, applicability_rate=0):

    """

    :param consumers: Consumer ids for which profile should be generated and set

    :type consumers: list

    :param size: A maximum number of packages in consumer profile

    :type size: int

    :param applicability_rate: N of applicable packages/N of packages in consumer profile

    :type applicability_rate: float

    """

    rpm_total = rcu_c.find({'unit_type_id': 'rpm'}).count()

    rcu_repo_ids = rcu_c.distinct('repo_id')

    distributors = []

    for repo_id in rcu_repo_ids:

        dist = distributors_c.find({'repo_id': repo_id}, projection=['distributor_id'])[0]

        distributors.append(dist['distributor_id'])

    _logger.info('all content units fetched')

    projection = ['name', 'epoch', 'version', 'release', 'arch', 'vendor']

    rpms = list(rpm_c.find({}, projection=projection))

    _logger.info('all rpm data fetched')

    applicable_rpms_idx = random.sample(range(size), int(size*applicability_rate))

    for consumer_id in consumers:

        wanted_rpms_idx = random.sample(range(rpm_total), size)

        profile = []

        for appl_idx, rpm_idx in enumerate(wanted_rpms_idx):

            pkg = {key: rpms[rpm_idx].get(key) for key in projection}

            if appl_idx in applicable_rpms_idx:

                pkg['version'] = '0.0.1'

            profile.append(pkg)

        set_consumer_profile(consumer_id, profile)

        for repo_id, dist_id in zip(rcu_repo_ids, distributors):

            bind_consumer(consumer_id, repo_id, dist_id)

if __name__ == '__main__':

    drop_consumer_collections()

    generate_and_set_consumers(500, 1000, 0.01)