ws-client.py

import argparse
import json
import sys
import time

from websocket import create_connection

RESULT_POLLING_FREQUENCY=60

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", "--target_branch", 
                        help="the target branch to fork and run the tests on")
    parser.add_argument("-c", "--target_commit", 
                        help="the target commit to checkout to run the tests on")
    parser.add_argument("-m", "--mode",
                        help="the execution mode. `run` runs and waits until the results are there, `wait` just waits, and `check` just returns the current task",
                        choices=['run', 'wait', 'check'],
                        default='run')
    args = parser.parse_args()
    return args

def issue_test_run(websocket, target_commit, target_branch):
    run_tests_req_data = {"cmd": {"job": "issue",
                                "benchmark": "CORRECTNESS",
                                "commit": target_commit,
                                "branch": target_branch,
                                }}
    msg = json.dumps(run_tests_req_data) 
    websocket.send(msg)
    print("POSIX Tests request made for branch:", target_branch, "and commit:", target_commit, file=sys.stderr)

def fetch_runs(websocket):
    data = {"cmd": {"job": "/fetch_runs", 
                    "count": 50}}
    msg = json.dumps(data) 
    # print("Sending:", msg, file=sys.stderr)
    websocket.send(msg)
    # print("Sent!", file=sys.stderr)
    res = websocket.recv()
    runs_data = json.loads(res)    
    return runs_data

def current_task(websocket):
    data = {"cmd": {"job": "/current_task"}}
    msg = json.dumps(data) 
    # print("Sending:", msg, file=sys.stderr)
    websocket.send(msg)
    # print("Sent!", file=sys.stderr)
    res = websocket.recv()
    res_data = json.loads(res)    
    return res_data

def wait_for_result(websocket, target_commit):
    found = False
    sleep_duration = RESULT_POLLING_FREQUENCY

    while not found:

        ## Fetch all runs
        runs_data = fetch_runs(websocket)
        result_rows = runs_data["data"]["rows"]

        ## Check if target commit exists in latest runs
        for row in result_rows:
            # print(row["commit"], row["bench"], file=sys.stderr)
            if row["commit"] == target_commit and row["bench"] == "CORRECTNESS":
                print("FOUND COMMIT:", target_commit, file=sys.stderr)
                found = True
                result_row = json.dumps(row)

        if not found:
            print("Results not present for commit:", target_commit, file=sys.stderr)
            print("Sleeping for", sleep_duration, "seconds", file=sys.stderr)
            time.sleep(sleep_duration)
    return result_row


## Parse the arguments
args = parse_args()

## Create a websocket connection
ws = create_connection("ws://antikythera.csail.mit.edu:52277/")

## Just check and exit
if args.mode == "check":
    ## Debug only
    print(current_task(ws))

elif args.mode in ["run", "wait"]:
    # target_commit = "e5937bca7644c6cae0a1b8214fa980b12145bbed"
    # target_branch = "pip-libdash"
    target_commit = args.target_commit
    target_branch = args.target_branch

    if args.mode == "run":
        ## Issue the POSIX tests requests
        issue_test_run(ws, target_commit, target_branch)
    
    ##
    ## Wait until we have the POSIX test results
    ##
    result_row = wait_for_result(ws, target_commit)

    ## Output the relevant test
    print(result_row)


ws.close()