1#! /usr/bin/python 2 3"""A simple heartbeat client. 4 5Executes heartbeats against a simple_heartbeat_server running on the give 6--server address and deserializes records into an in memory sqlite database. 7 8Usage: 91. heartbeat_client.py 10 --server http://localhost:8080 11 --board lumpy 12 13 Perform a heartbeat against the given server for the given board, 14 and deserialize records into a sqlite database. 15 162. heartbeat_client.py 17 --server http://localhost:8080 18 --board lumpy 19 --host_limit 1 --job_limit 100 20 21 Do the same as 1, but instruct the server to limit the hosts to 1 22 and jobs to 100. This is useful for debugging issues with only jobs/hosts. 23""" 24 25 26from json import decoder 27import argparse 28import sys 29import urllib2 30 31import common 32 33from autotest_lib.scheduler.shard import simple_heartbeat_server 34from autotest_lib.frontend import setup_django_environment 35from autotest_lib.frontend.afe import frontend_test_utils 36from autotest_lib.frontend.afe import models 37 38json_decoder = decoder.JSONDecoder() 39 40 41class HeartbeatHandler(frontend_test_utils.FrontendTestMixin): 42 """Performs heartbeats and deserializes into an in memory database.""" 43 44 _config_section = 'AUTOTEST_WEB' 45 46 47 def __init__(self, server): 48 """Initialize a heartbeat server. 49 50 @param server: The address of a simple_heartbeat_server. 51 """ 52 self.server = server 53 self._frontend_common_setup(setup_tables=True, fill_data=False) 54 55 56 @staticmethod 57 @simple_heartbeat_server.time_call 58 def get_heartbeat_packet(server, board, host_limit, job_limit): 59 """Perform the heartbeat. 60 61 Constructs a url like: http://localhost:8080/lumpy?raw&host_limit=3 62 and does a urlopen. 63 64 @param server: The address of a simple_heartbeat_server. 65 @param host_limit: The number of hosts to include in the heartbeat. 66 @param job_limit: The number of jobs to include in the heartbeat. 67 68 @return: A string containing the heartbeat packet. 69 """ 70 url = '%s/%s?raw' % (server, board) 71 if job_limit: 72 url = '%s&job_limit=%s' % (url, job_limit) 73 if host_limit: 74 url = '%s&host_limit=%s' % (url, host_limit) 75 print 'Performing heartbeat against %s' % url 76 return urllib2.urlopen(url).read() 77 78 79 @staticmethod 80 @simple_heartbeat_server.time_call 81 def deserialize_heartbeat(packet): 82 """Deserialize the given heartbeat packet into an in memory database. 83 84 @param packet: A string representing the heartbeat packet containing 85 jobs and hosts. 86 87 @return: The json decoded heartbeat. 88 """ 89 response = json_decoder.decode(packet) 90 [models.Host.deserialize(h) for h in response['hosts']] 91 [models.Job.deserialize(j) for j in response['jobs']] 92 return response 93 94 95 def perform_heartbeat(self, board, host_limit, job_limit): 96 """Perform a heartbeat against the given server, for the given board. 97 98 @param board: Boardname, eg: lumpy. 99 @param host_limit: Limit number of hosts retrieved. 100 @param job_limit: Limit number of jobs retrieved. 101 """ 102 timing, packet = self.get_heartbeat_packet( 103 self.server, board, host_limit, job_limit) 104 print 'Time to perform heartbeat %s' % timing 105 timing, response = self.deserialize_heartbeat(packet) 106 print 'Time to deserialize hearbeat %s' % timing 107 print ('Jobs: %s, Hosts: %s' % 108 (len(response['jobs']), len(response['hosts']))) 109 110 111def _parse_args(args): 112 parser = argparse.ArgumentParser( 113 description='Start up a simple heartbeat client.') 114 parser.add_argument( 115 '--server', default='http://localhost:8080', 116 help='Address of a simple_heartbeat_server to heartbeat against.') 117 parser.add_argument( 118 '--board', default='lumpy', 119 help='Heartbeats can only be performed ' 120 'against a specific board, eg: lumpy.') 121 parser.add_argument( 122 '--host_limit', default='', 123 help='Limit hosts in the heartbeat.') 124 parser.add_argument( 125 '--job_limit', default='', 126 help='Limit jobs in the heartbeat.') 127 args = parser.parse_args(args) 128 args.board = args.board.lstrip( 129 simple_heartbeat_server.BoardHandler.board_prefix) 130 return args 131 132 133if __name__ == '__main__': 134 args = _parse_args(sys.argv[1:]) 135 HeartbeatHandler(args.server).perform_heartbeat( 136 args.board, args.host_limit, args.job_limit) 137