sched

#!/usr/bin/env python
import sys, string, urllib
import requests
from requests.packages.urllib3.util.retry import Retry
from requests.adapters import HTTPAdapter
import optparse

session = requests.Session()
adapter = HTTPAdapter(
  max_retries=Retry(
    connect=5,
    status=5,
    backoff_factor=0.1,
    status_forcelist=[500, 502, 503, 504]
  )
)
session.mount('http://', adapter)
session.mount('https://', adapter)


def test_time(target, test_name, runtime):
  r = session.post(target + "/record/%s/%f" % (urllib.quote(test_name, safe=""), runtime))
  print r.text.encode('utf-8')
  assert r.status_code == 204

def test_sched(target, test_run, shard_count, shard_id):
  tests = {'tests': string.split(sys.stdin.read())}
  r = session.post(target + "/schedule/%s/%d/%d" % (test_run, shard_count, shard_id), json=tests)
  assert r.status_code == 200
  result = r.json()
  for test in sorted(result['tests']):
    print test.encode('utf-8')

def usage():
  print "%s (--target=...) <cmd> <args..>" % sys.argv[0]
  print "   time <test name> <run time>"
  print "   sched <test run> <num shards> <shard id>"

def main():
  parser = optparse.OptionParser()
  parser.add_option('--target', default="http://positive-cocoa-90213.appspot.com")
  options, args = parser.parse_args()
  if len(args) < 3:
    usage()
    sys.exit(1)

  if args[0] == "time":
    test_time(options.target, args[1], float(args[2]))
  elif args[0] == "sched":
    test_sched(options.target, args[1], int(args[2]), int(args[3]))
  else:
    usage()

if __name__ == '__main__':
  main()