1 files changed, 39 insertions(+), 13 deletions(-)

M main.go
M main.go +39 -13
@@ 154,10 154,18 @@ func CSVWriter(output io.Writer, workers
 	return c
 }
 
-func main() {
+//
+// Exits program if there's any error parsing the arguments
+//
+func parse_arguments() (
+	url string,
+	capabilities selenium.Capabilities,
+	workers int,
+) {
 	var args struct {
-		Workers uint `short:"w" long:"workers" default:"4" description:"Number of workers"`
-		Arg     struct {
+		Workers      uint              `short:"w" default:"4" description:"Number of workers"`
+		Capabilities map[string]string `short:"c" default:"browserName:firefox" description:"Selenium capabilities"`
+		Arg          struct {
 			URL string `description:"URL to the Selenium server"`
 		} `positional-args:"yes" required:"yes"`
 	}

          
@@ 173,24 181,42 @@ func main() {
 		logger.Fatalln("Extra arguments:", extra)
 	}
 
-	caps := selenium.Capabilities{
-		"idleTimeout": 1000,
-		"browserName": "firefox",
-		"platform":    "linux",
+	url = args.Arg.URL
+	workers = int(args.Workers)
+	capabilities = make(selenium.Capabilities, len(args.Capabilities))
+	for key, value := range args.Capabilities {
+		capabilities[key] = value
 	}
-	var workers = int(args.Workers)
+
+	return
+}
+
+func main() {
+	var url, capabilities, workers = parse_arguments()
+
 	var input = ScanLines(os.Stdin)
-	var output = CSVWriter(os.Stdout, workers)
-	var url = args.Arg.URL
+	// Buffer results to avoid blocking workers
+	var output = make(chan Result, workers)
 	var wg sync.WaitGroup
 
 	logger.Println("Using", workers, "workers with", url)
 
 	wg.Add(int(workers))
-	for i := 0; i < workers; i++ {
-		go worker(caps, url, &wg, input, output)
+	for i := 0; i < int(workers); i++ {
+		go worker(capabilities, url, &wg, input, output)
+	}
+
+	var writer = csv.NewWriter(os.Stdout)
+
+	for o := range output {
+		var e = ""
+		if o.Error != nil {
+			logger.Println("Error for", o.URL, ":", o.Error)
+			e = o.Error.Error()
+		}
+		writer.Write([]string{o.URL, e, o.Duration.String()})
+		writer.Flush()
 	}
 
 	wg.Wait()
-	close(output)
 }