check_pending_for_worker_type.py

import argparse
import json
import sys
import urllib2

QUEUE_PENDING_URL = "https://queue.taskcluster.net/v1/pending/{}/{}"
STATUS_CODE = {'OK': 0, 'WARNING': 1, "CRITICAL": 2, "UNKNOWN": 3}

def get_pending_count_for_worker_type(provisioner, worker_type):
    pending_url = QUEUE_PENDING_URL.format(provisioner, worker_type)
    response = urllib2.urlopen(pending_url, timeout=30)
    result = json.loads(response.read())
    return result['pendingTasks']

def pending_status(pending, critical_threshold, warning_threshold):
    if pending >= critical_threshold:
        return 'CRITICAL'
    elif pending >= warning_threshold:
        return 'WARNING'
    else:
        return 'OK'

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--provisioner', action='store', dest='provisioner',
        required=True, metavar="PROVISIONER", help='Specify provisioner ID for given worker type')
    parser.add_argument(
        '--workertype', action='store', dest='worker_type',
        required=True, metavar="WORKERTYPE", help='Specify worker type to query pending count for')
    parser.add_argument(
        '-c', '--critical', action='store', type=int, dest='critical_threshold',
        default=2000, metavar="CRITICAL", help='Set CRITICAL level as integer eg. 2000')
    parser.add_argument(
        '-w', '--warning', action='store', type=int, dest='warning_threshold',
        default=1200, metavar="WARNING", help='Set WARNING level as integer eg. 1200')

    args = parser.parse_args()
    try:
        pending = get_pending_count_for_worker_type(args.provisioner, args.worker_type)
        status = pending_status(pending, args.critical_threshold, args.warning_threshold)
        print "{} Pending Jobs: {} for {}".format(status, pending, args.worker_type)
        sys.exit(STATUS_CODE[status])
    except Exception as e:
        print e
        sys.exit(STATUS_CODE['UNKNOWN'])