Allow configuring up and down request sizes from flags

diff --git a/main.go b/main.go
index 0ca426f..1039d2f 100644 (file)
--- a/main.go
+++ b/main.go
@@ -3,6 +3,7 @@ package main
 import (
        "bufio"
        "bytes"
+       "flag"
        "fmt"
        "io"
        "log"
@@ -48,8 +49,12 @@ const DOWNLOAD_PATH = "https://speed.cloudflare.com/__down"
 var iter = 10
 
 func main() {
-       bytes := 10 * MiB
-       dUrl, _ := url.Parse(DOWNLOAD_PATH + "?bytes=" + strconv.FormatInt(int64(bytes), 10))
+       upFlag := flag.Int64("up", 10, "number of MB for upload test")
+       downFlag := flag.Int64("down", 1, "number of MB for download test")
+       flag.Parse()
+       upBytes := Byte(*upFlag) * MiB
+       downBytes := Byte(*downFlag) * MiB
+       dUrl, _ := url.Parse(DOWNLOAD_PATH + "?bytes=" + strconv.FormatInt(int64(downBytes), 10))
        uUrl, _ := url.Parse(UPLOAD_PATH)
        down := &http.Request{
                Method: "GET",
@@ -62,8 +67,8 @@ func main() {
        upChan := make(chan float64)
 
        client := &http.Client{}
-       go measureNRequests(iter, client, func() *http.Request { return down }, bytes, downChan)
-       go measureNRequests(iter, client, func() *http.Request { return getUp(uUrl, bytes) }, bytes, upChan)
+       go measureNRequests(iter, client, func() *http.Request { return down }, downBytes, downChan)
+       go measureNRequests(iter, client, func() *http.Request { return getUp(uUrl, upBytes) }, upBytes, upChan)
        tickChan := make(chan rune)
        go func() {
                for {