From 99373106d502b54566e6d85aa20a3b6dd984f181 Mon Sep 17 00:00:00 2001
From: Mihir Joshi <mihir67mj@gmail.com>
Date: Sat, 13 Jan 2024 23:45:21 +0530
Subject: [PATCH] Add HTTP Methods support, Request Body Support

- Set default header as Content-Type:application/json
---
 load_tester.go | 14 ++++++++++----
 main.go        | 17 +++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/load_tester.go b/load_tester.go
index 3f0c76c..6e4fed3 100644
--- a/load_tester.go
+++ b/load_tester.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/http/httptrace"
 	"os"
+	"strings"
 	"sync/atomic"
 	"time"
 )
@@ -15,6 +16,8 @@ var (
 	TotalReq   int
 	Endpoint   string
 	Concurrent int
+	HttpMethod string
+	Body       string
 
 	// Accessed by other files to show results
 	ReqProgress                   int
@@ -33,9 +36,10 @@ var (
 )
 
 const (
-	failed       = "failed"
-	succeeded    = "succeeded"
-	reqPerSecond = "reqPerSecond"
+	failed        = "failed"
+	succeeded     = "succeeded"
+	reqPerSecond  = "reqPerSecond"
+	totalDuration = "totalDuration"
 )
 
 type Response struct {
@@ -62,7 +66,8 @@ func LoadTest() {
 func createRequestJobs(reqPool chan<- *http.Request, url string, numberOfRequests int) {
 	defer close(reqPool)
 	for i := 0; i < numberOfRequests; i++ {
-		r, err := http.NewRequest(http.MethodGet, url, nil)
+		r, err := http.NewRequest(HttpMethod, url, strings.NewReader(Body))
+		r.Header.Set("Content-Type", "application/json")
 		if err != nil {
 			panic(err)
 		}
@@ -96,6 +101,7 @@ func evaluateResponses(responseChannel <-chan *Response) {
 	results[failed] = fmt.Sprintf("%d", failedCount)
 	requestsPerSecond := float64(succeededCount) / Elapsed.Seconds()
 	results[reqPerSecond] = fmt.Sprintf("%f", requestsPerSecond)
+	results[totalDuration] = Elapsed.String()
 }
 
 func startRequestWorkers(requestChannel <-chan *http.Request, responseChannel chan<- *Response, maxConcurrentRequests int) {
diff --git a/main.go b/main.go
index 69abc19..466ba13 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"github.com/fatih/color"
+	"net/http"
 	url2 "net/url"
 	"os"
 	"runtime"
@@ -18,6 +19,8 @@ func main() {
 	flag.StringVar(&Endpoint, "url", "", "Endpoint URL for load testing")
 	flag.IntVar(&TotalReq, "n", defaultNumberOfTotalRequests, "Total number of requests to make")
 	flag.IntVar(&Concurrent, "c", defaultConcurrentRequests, "Number of Concurrent requests")
+	flag.StringVar(&HttpMethod, "method", "", "HTTP Method to use while making the request")
+	flag.StringVar(&Body, "body", "", "JSON Request Body for each Request")
 	flag.Parse()
 	if _, err := url2.ParseRequestURI(Endpoint); err != nil {
 		fmt.Printf("Invalid Endpoint URL: %s\n", err)
@@ -25,8 +28,21 @@ func main() {
 		os.Exit(-1)
 	}
 
+	if HttpMethod == "" {
+		HttpMethod = http.MethodGet
+	}
+
+	switch HttpMethod {
+	case http.MethodGet, http.MethodHead, http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete, http.MethodConnect, http.MethodOptions, http.MethodTrace:
+	default:
+		fmt.Printf("Invalid Http Method: %s\n", HttpMethod)
+		os.Exit(-1)
+	}
+
 	println("USING:", runtime.NumCPU(), "CPUs")
 	println("URL:", Endpoint)
+	println("HTTP Method:", HttpMethod)
+	println("Request Body:", Body)
 	println("Total number of requests:", TotalReq)
 	println("Parallel requests:", Concurrent)
 
@@ -40,4 +56,5 @@ func printResults() {
 	color.Green("Succeeded Requests: %s", results[succeeded])
 	color.Red("Failed Requests: %s", results[failed])
 	color.Cyan("Requests/Second: %s", results[reqPerSecond])
+	color.Cyan("Completed Load Testing In: %s", results[totalDuration])
 }