1 # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 2 # Use of this source code is governed by a BSD-style license that can be 3 # found in the LICENSE file. 4 5 """Helper functions for testing stats module and elasticsearch 6 """ 7 8 import logging 9 import time 10 11 import common 12 13 import elasticsearch 14 15 from autotest_lib.client.common_lib.cros.graphite import es_utils 16 from autotest_lib.client.common_lib.cros.graphite import autotest_stats 17 18 19 # Defines methods in the stats class that can take in metadata. 20 TARGET_TO_STATS_CLASS = { 21 'timer': autotest_stats.Timer, 22 'gauge': autotest_stats.Gauge, 23 'raw': autotest_stats.Raw, 24 'average': autotest_stats.Average, 25 'counter': autotest_stats.Counter, 26 } 27 28 # Maps target type to method to trigger sending of metadata. 29 # This differs based on what each object does. 30 # For example, in order for timer to send something, its stop 31 # method must be called. This differs for other stats objects. 32 TARGET_TO_METHOD = { 33 'timer': 'stop', 34 'gauge': 'send', 35 'raw': 'send', 36 'average': 'send', 37 'counter': '_send', 38 } 39 40 # Default maximum number of entries to return from ES query 41 DEFAULT_NUM_ENTRIES = 100 42 43 class EsTestUtilException(Exception): 44 """Exception raised when functions here fail. """ 45 pass 46 47 48 def sequential_random_insert_ints(keys, num_entries, target_type, index, 49 host, port, use_http, udp_port, 50 between_insert_secs=0, 51 print_interval=10): 52 """Inserts a bunch of random entries into the es database. 53 Keys are given, values are randomly generated. 54 55 @param keys: A list of keys 56 @param num_entries: Number of entries to insert 57 @param target_type: This must be in 58 ['timer', 'gauge', 'raw', 'average', 'counter'] 59 @param between_insert_secs: Time to sleep after each insert. 60 defaults to no sleep time. 61 @param print_interval: how often to print 62 defaults to every 10 entries. 63 @param index: Index of es db to insert to 64 @param host: host of es db 65 @param port: port of es db 66 """ 67 # We are going to start the value at 0 and increment it by one per val. 68 for i in range(num_entries): 69 if print_interval == 0 or i % print_interval == 0: 70 print(' Inserting entry #%s with keys %s into index "%s."' 71 % (i, str(keys), index)) 72 metadata = {} 73 for value, key in enumerate(keys): 74 metadata[key] = value 75 76 # Subname and value are not important from metadata pov. 77 subname = 'metadata.test' 78 value = 10 79 stats_target = TARGET_TO_STATS_CLASS[target_type](subname, 80 metadata=metadata, 81 es=es_utils.ESMetadata(use_http=use_http, host=host, 82 port=port, index=index, 83 udp_port=udp_port)) 84 85 if target_type == 'timer': 86 stats_target.start() 87 stats_target.stop() 88 else: 89 getattr(stats_target, TARGET_TO_METHOD[target_type])(subname, value) 90 time.sleep(between_insert_secs) 91 92 93 def clear_index(index, host, port, timeout, sleep_time=0.5, clear_timeout=5): 94 """Clears index in es db located at host:port. 95 96 Warning: Will delete all data in es for a given index 97 98 @param index: Index of es db to clear 99 @param host: elasticsearch host 100 @param port: elasticsearch port 101 @param timeout: how long to wait while connecting to es. 102 @param sleep_time: time between tries of clear_index 103 defaults to 0.5 seconds 104 @param clear_timeout: how long to wait for index to be cleared. 105 defualts to 5 seconds 106 Will quit and throw error if not cleared. (Number of seconds) 107 """ 108 es = elasticsearch.Elasticsearch(host=host, 109 port=port, 110 timeout=timeout) 111 if es.indices.exists(index=index): 112 print 'deleting index %s' % (index) 113 es.indices.delete(index=index) 114 time_start = time.time() 115 while es.indices.exists(index=index): 116 print 'waiting until index is deleted...' 117 time.sleep(sleep_time) 118 if time.time() - time_start > clear_timeout: 119 raise EsTestUtilException('clear_index failed.') 120 121 print 'successfully deleted index %s' % (index) 122
