allow other HTTP methods

Author: mfbmina

cmd/root.go

@@ -1,19 +1,9 @@
 package cmd
 
 import (
-	"fmt"
-
-	"github.com/getsentry/sentry-go"
-	"github.com/mfbmina/enxame/reporter"
-	"github.com/mfbmina/enxame/swarm"
 	"github.com/spf13/cobra"
 )
 
-var format string
-var output string
-var requestsPerUser int
-var users int
-
 var rootCmd = &cobra.Command{
 	Use:   "enxame",
 	Short: "A http load tester and benchmarking utility made in Go",
@@ -29,34 +19,3 @@ func init() {
 
 	rootCmd.AddCommand(runCmd())
 }
-
-func runCmd() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "run",
-		Short: "Run a test",
-		Long:  "Run a test against the given URI",
-		Args:  cobra.MinimumNArgs(1),
-		Run: func(cmd *cobra.Command, args []string) {
-			fmt.Println("Swarming", args[0], "...")
-			responses := swarm.Swarm(args[0], requestsPerUser, users)
-
-			fmt.Printf("Reporting results as %s...\n", format)
-			r := reporter.NewReporter(format, output, responses)
-			feedback, err := r.Report()
-			if err != nil {
-				sentry.CaptureException(err)
-				return
-			}
-
-			fmt.Println(feedback)
-			fmt.Println("Done!")
-		},
-	}
-
-	cmd.Flags().IntVarP(&requestsPerUser, "requests_per_user", "r", 5, "Max number of requests per user")
-	cmd.Flags().IntVarP(&users, "users", "u", 10, "Max number of concurrent users")
-	cmd.Flags().StringVarP(&format, "format", "f", "txt", "The report format (txt, csv, json)")
-	cmd.Flags().StringVarP(&output, "output", "o", "", "The output file name (extension will be appended automatically)")
-
-	return cmd
-}

cmd/run.go

@@ -0,0 +1,78 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/getsentry/sentry-go"
+	"github.com/mfbmina/enxame/reporter"
+	"github.com/mfbmina/enxame/swarm"
+	"github.com/spf13/cobra"
+)
+
+var format string
+var method string
+var output string
+var requestsPerUser int
+var users int
+
+func runCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "run",
+		Short: "Run a test",
+		Long:  "Run a test against the given URI",
+		Args:  cobra.MinimumNArgs(1),
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			fmt.Println("Validating method...")
+			fmt.Println("Method:", method)
+			if !validMethod(method) {
+				err := fmt.Errorf("Invalid method. Please use one of the following: GET, POST, PUT, DELETE, PATCH, OPTIONS, HEAD")
+				sentry.CaptureException(err)
+				return err
+			}
+
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			fmt.Println("Swarming", args[0], "...")
+			responses, err := swarm.Swarm(args[0], method, requestsPerUser, users)
+			if err != nil {
+				sentry.CaptureException(err)
+				return err
+			}
+
+			fmt.Printf("Reporting results as %s...\n", format)
+			r := reporter.NewReporter(format, output, responses)
+			feedback, err := r.Report()
+			if err != nil {
+				sentry.CaptureException(err)
+				return err
+			}
+
+			fmt.Println(feedback)
+			fmt.Println("Done!")
+			return nil
+		},
+	}
+
+	cmd.Flags().IntVarP(&requestsPerUser, "requests_per_user", "r", 5, "Max number of requests per user")
+	cmd.Flags().IntVarP(&users, "users", "u", 10, "Max number of concurrent users")
+	cmd.Flags().StringVarP(&format, "format", "f", "txt", "The report format (txt, csv, json)")
+	cmd.Flags().StringVarP(&method, "method", "X", "GET", "The http method to use (GET, POST, PUT, DELETE, PATCH, OPTIONS, HEAD)")
+	cmd.Flags().StringVarP(&output, "output", "o", "", "The output file name (extension will be appended automatically)")
+
+	return cmd
+}
+
+func validMethod(method string) bool {
+	allowedMethods := []string{"GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"}
+	f := false
+
+	for _, m := range allowedMethods {
+		if m == method {
+			f = true
+			break
+		}
+	}
+
+	return f
+}

swarm/request.go

@@ -11,33 +11,37 @@ type HTTPResponse struct {
 	Path       string        `json:"path"`
 }
 
-func Swarm(path string, requestsPerUser int, concurrentUsers int) []HTTPResponse {
+func Swarm(path, method string, requestsPerUser, concurrentUsers int) ([]HTTPResponse, error) {
 	totalRequests := requestsPerUser * concurrentUsers
 	channel := make(chan HTTPResponse, totalRequests)
 	responses := make([]HTTPResponse, totalRequests)
 
 	for i := 0; i < concurrentUsers; i++ {
-		go userRequests(path, requestsPerUser, channel)
+		req, _ := http.NewRequest(method, path, nil)
+		go userRequests(req, requestsPerUser, channel)
 	}
 
 	for i := 0; i < totalRequests; i++ {
 		responses[i] = <-channel
 	}
 
-	return responses
+	return responses, nil
 }
 
-func userRequests(path string, requestsPerUser int, channel chan HTTPResponse) {
+func userRequests(req *http.Request, requestsPerUser int, channel chan HTTPResponse) {
 	for i := 0; i < requestsPerUser; i++ {
 		startTime := time.Now()
-		resp, err := http.Get(path)
+
+		client := &http.Client{}
+		resp, err := client.Do(req)
+
 		elapsedTime := time.Since(startTime) / time.Millisecond
 
 		if err != nil {
-			channel <- HTTPResponse{StatusCode: 0, Path: path, Time: elapsedTime}
+			channel <- HTTPResponse{StatusCode: 0, Path: req.URL.String(), Time: elapsedTime}
 			continue
 		}
 
-		channel <- HTTPResponse{StatusCode: resp.StatusCode, Path: path, Time: elapsedTime}
+		channel <- HTTPResponse{StatusCode: resp.StatusCode, Path: req.URL.String(), Time: elapsedTime}
 	}
 }

swarm/request_test.go

@@ -8,7 +8,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestSwarmWhenApiCallisSuccessful(t *testing.T) {
+func TestSwarmWhenApiCallWithGet(t *testing.T) {
 	url := "http://example.com/users"
 
 	mock := apitest.NewMock().
@@ -18,24 +18,44 @@ func TestSwarmWhenApiCallisSuccessful(t *testing.T) {
 
 	defer mock.EndStandalone()()
 
-	r := Swarm(url, 1, 1)
+	r, err := Swarm(url, "GET", 1, 1)
 	assert.Equal(t, 1, len(r))
 	assert.Equal(t, 200, r[0].StatusCode)
 	assert.NotNil(t, r[0].Time)
+	assert.Nil(t, err)
+}
+
+func TestSwarmWhenApiCallWithPost(t *testing.T) {
+	url := "http://example.com/users"
+
+	mock := apitest.NewMock().
+		Post(url).
+		RespondWith().
+		Status(http.StatusCreated)
+
+	defer mock.EndStandalone()()
+
+	r, err := Swarm(url, "POST", 1, 1)
+	assert.Equal(t, 1, len(r))
+	assert.Equal(t, 201, r[0].StatusCode)
+	assert.NotNil(t, r[0].Time)
+	assert.Nil(t, err)
 }
 
 func TestSwarmWhenApiCallReturnsError(t *testing.T) {
 	url := "strange_url"
 
-	r := Swarm(url, 1, 1)
+	r, err := Swarm(url, "GET", 1, 1)
 	assert.Equal(t, 1, len(r))
 	assert.Equal(t, 0, r[0].StatusCode)
 	assert.NotNil(t, r[0].Time)
+	assert.Nil(t, err)
 }
 
 func TestSwarmReturnsTotalRequests(t *testing.T) {
 	url := "strange_url"
 
-	r := Swarm(url, 2, 5)
+	r, err := Swarm(url, "GET", 2, 5)
 	assert.Equal(t, 10, len(r))
+	assert.Nil(t, err)
 }