From 2e062a689915aac4246d7449d9d0be4a3da36baf Mon Sep 17 00:00:00 2001
From: Aleksandr Razumov <ernado@go-faster.org>
Date: Sun, 7 Jan 2024 17:01:43 +0300
Subject: [PATCH] fix(promrw): fix target parsing

---
 cmd/promrw/bench.go | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/cmd/promrw/bench.go b/cmd/promrw/bench.go
index 6b910d85..1fed9edb 100644
--- a/cmd/promrw/bench.go
+++ b/cmd/promrw/bench.go
@@ -5,7 +5,6 @@ import (
 	"bytes"
 	"context"
 	"crypto/sha256"
-	"flag"
 	"fmt"
 	"io"
 	"math/rand"
@@ -495,19 +494,18 @@ func (s *Bench) prometheusConfig() *config {
 	return cfg
 }
 
-func (s *Bench) parseTargets() {
-	for _, arg := range flag.Args() {
+func (s *Bench) parseTargets(args []string) error {
+	for _, arg := range args {
 		u, err := url.Parse(arg)
 		if err != nil {
-			fmt.Fprintln(os.Stderr, "invalid target:", err)
-			os.Exit(1)
+			return errors.Wrap(err, "parse url")
 		}
 		s.targets = append(s.targets, u.String())
 	}
 	if len(s.targets) == 0 {
-		fmt.Fprintln(os.Stderr, "no targets specified")
-		os.Exit(1)
+		return errors.New("no targets")
 	}
+	return nil
 }
 
 func (s *Bench) waitForTarget(ctx context.Context, target string) error {
@@ -697,8 +695,11 @@ func newBenchCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "bench",
 		Short: "Start remote write benchmark",
+		Args:  cobra.MinimumNArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			b.parseTargets()
+			if err := b.parseTargets(args); err != nil {
+				return err
+			}
 			return b.run(cmd.Context())
 		},
 	}