test_my_server

#!/usr/bin/env python3
"""
Usage: ./test_my_server -c(cpu)/-i(io) -nc? <clients> <host> <port>

A generic script to test an HTTP server's performance
"""

import requests
import argparse
import multiprocessing
import time

# of parallel requests
CLIENTS = 5

def send_http_request(host, port, endpoint):
    url = f"http://{host}:{port}/{endpoint}"
    start = time.time()
    response = requests.get(url)
    end = time.time()
    print(f"{url} - Status Code: {response.status_code}, took {(end-start)*1000:.2f} ms")
    #print(response.text)
    #print("\n")


def send_parallel_requests(host, port, endpoint):
    processes = []
    for i in range(CLIENTS):
        p = multiprocessing.Process(target=send_http_request, args=(host, port, endpoint))
        processes.append(p)
    start = time.time()
    for process in processes:
        process.start()

    for process in processes:
        process.join()

    end = time.time()
    # Note that this includes time of process creation
    # and termination. Assuming it is miniscule
    # compared to time taken by http requests.
    print(f"Total time taken {(end-start)*1000:.2f} ms")


def get_args():
    parser = argparse.ArgumentParser(description="A script to test performance of a generic server") 
    parser.add_argument("-c", "--cpu", action='store_true',
                         help="Flag to test performance of cpu bound tasks")
    parser.add_argument("-i", "--io" , action='store_true',
                         help="Flag to test performance of i/o bound tasks")
    parser.add_argument('-nc', '--clients', type=int, help='Number of clients')
    parser.add_argument("host")
    parser.add_argument("port") 
    args = parser.parse_args()
    return args


def main():
    args = get_args()
    endpoint = ""
    if args.cpu:
        endpoint = "cpu"
    if args.io:
        endpoint = "io"
    if args.cpu and args.io:
        raise ValueError("Can't test cpu and io performance together. Please use -c OR -i")
    if not (args.cpu or args.io):
        raise ValueError("Specify cpu and io performance to test. Please use -c OR -i")
    if args.clients:
        global CLIENTS
        CLIENTS = args.clients

    send_parallel_requests(args.host, args.port, endpoint)
    

if __name__ == '__main__':
    main()