Project

Profile

Help

Task #2449 » gen_test_data.py

To generate consumer data with various applicability rate - ttereshc, 07/04/2017 12:14 AM

 
1
import logging
2
import logging.handlers
3
import random
4
import uuid
5

    
6
from pulp.server import initialization
7
from pulp.server.db import connection
8
from pulp.server.controllers import consumer as consumer_controller
9
from pulp.server.managers import factory
10

    
11
_logger = logging.getLogger(__name__)
12
_logger.setLevel(logging.INFO)
13
_logger.addHandler(logging.handlers.SysLogHandler(address='/dev/log'))
14

    
15

    
16
initialization.initialize()
17

    
18
db = connection.get_database()
19
consumers_c = db['consumers']
20
bindings_c = db['consumer_bindings']
21
profiles_c = db['consumer_unit_profiles']
22
repos_c = db['repos']
23
rpa_c = db['repo_profile_applicability']
24
rcu_c = db['repo_content_units']
25
rpm_c = db['units_rpm']
26
erratum_c = db['units_erratum']
27
distributors_c = db['repo_distributors']
28

    
29

    
30
def drop_consumer_collections():
31
    consumers_c.drop()
32
    bindings_c.drop()
33
    profiles_c.drop()
34
    rpa_c.drop()
35

    
36

    
37
def generate_and_set_consumers(num=500, profile_size=1000, applicability_rate=0):
38
    """
39
    :param num: number of consumer profiles to generate
40
    :type num: int
41
    :param profile_size: number of packages in consumer profile
42
    :type profile_size: int
43
    :param applicability_rate: N of applicable packages/N of packages in consumer profile
44
    :type applicability_rate: float
45
    """
46
    consumers = []
47
    for _ in range(num):
48
        consumer_id = str(uuid.uuid4())
49
        _logger.info('Registering consumer %s' % consumer_id)
50
        register_consumer(consumer_id)
51
        consumers.append(consumer_id)
52
    generate_consumer_profile(consumers, profile_size, applicability_rate)
53

    
54

    
55
def register_consumer(consumer_id):
56
    """
57
    :param consumer_id: Id of a consumer to register with
58
    :type consumer_id: str
59
    """
60
    manager = factory.consumer_manager()
61
    consumer, certificate = manager.register(consumer_id)
62

    
63

    
64
def set_consumer_profile(consumer_id, profile, content_type='rpm'):
65
    """
66
    :param consumer_id: Id of a consumer to set a profile for
67
    :type consumer_id: str
68
    :param profile:  A profile to associate to consumer
69
    :type profile: list
70
    """
71
    manager = factory.consumer_profile_manager()
72
    new_profile = manager.create(consumer_id, content_type, profile)
73
    _logger.info('Profile set')
74

    
75

    
76
def bind_consumer(consumer_id, repo_id, distributor_id):
77
    """
78
    :param consumer_id: Id of a consumer
79
    :type consumer_id: str
80
    :param repo_id: Id of repo to bind to
81
    :type repo_id: str
82
    :param distributor_id: Id of repo distributor to bind to
83
    :type distributor_id: str
84
    """
85
    notify_agent = False
86
    consumer_controller.bind(consumer_id, repo_id, distributor_id, notify_agent, {}, {})
87
    _logger.info('Consumer bound to %s' % repo_id)
88

    
89

    
90
def generate_consumer_profile(consumers, size, applicability_rate=0):
91
    """
92
    :param consumers: Consumer ids for which profile should be generated and set
93
    :type consumers: list
94
    :param size: A maximum number of packages in consumer profile
95
    :type size: int
96
    :param applicability_rate: N of applicable packages/N of packages in consumer profile
97
    :type applicability_rate: float
98
    """
99
    rpm_total = rcu_c.find({'unit_type_id': 'rpm'}).count()
100
    rcu_repo_ids = rcu_c.distinct('repo_id')
101
    distributors = []
102
    for repo_id in rcu_repo_ids:
103
        dist = distributors_c.find({'repo_id': repo_id}, projection=['distributor_id'])[0]
104
        distributors.append(dist['distributor_id'])
105

    
106
    _logger.info('all content units fetched')
107
    projection = ['name', 'epoch', 'version', 'release', 'arch', 'vendor']
108
    rpms = list(rpm_c.find({}, projection=projection))
109
    _logger.info('all rpm data fetched')
110

    
111
    applicable_rpms_idx = random.sample(range(size), int(size*applicability_rate))
112
    for consumer_id in consumers:
113
        wanted_rpms_idx = random.sample(range(rpm_total), size)
114
        profile = []
115
        for appl_idx, rpm_idx in enumerate(wanted_rpms_idx):
116
            pkg = {key: rpms[rpm_idx].get(key) for key in projection}
117
            if appl_idx in applicable_rpms_idx:
118
                pkg['version'] = '0.0.1'
119
            profile.append(pkg)
120
        set_consumer_profile(consumer_id, profile)
121
        for repo_id, dist_id in zip(rcu_repo_ids, distributors):
122
            bind_consumer(consumer_id, repo_id, dist_id)
123

    
124

    
125
if __name__ == '__main__':
126
    drop_consumer_collections()
127
    generate_and_set_consumers(500, 1000, 0.01)
(2-2/2)