Small changes

Shopify · Sep 17, 2021 · 18d56b9 · 18d56b9
1 parent d062d0d
commit 18d56b9
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 78 deletions.
diff --git a/cli/cli.go b/cli/cli.go
@@ -51,14 +51,14 @@ var toxicDescription = `
 
   toxic add:
     usage: toxiproxy-cli toxic add --type <toxicType> [--downstream|--upstream] \
-            --toxicName <toxicName> --toxicity <toxicity> \
-            --attribute <key=value> [--attribute <key2=value2>] <proxyName>\
+            --toxicName <toxicName> [--toxicity <float>] \
+            --attribute <key=value> [--attribute <key2=value2>] <proxyName>
 
 
     example: toxiproxy-cli toxic add -t latency -n myToxic -a latency=100 -a jitter=50 myProxy
 
   toxic update:
-    usage: toxiproxy-cli toxic update --toxicName <toxicName> \
+    usage: toxiproxy-cli toxic update --toxicName <toxicName> [--toxicity <float>] \
             --attribute <key1=value1> [--attribute <key2=value2>] <proxyName>
 
     example: toxiproxy-cli toxic update -n myToxic -a jitter=25 myProxy
@@ -447,54 +447,6 @@ func parseToxicity(c *cli.Context, defaultToxicity float32) (float32, error) {
 	return toxicity, nil
 }
 
-func parseToxicCommonParams(context *cli.Context) (*toxiproxy.ToxicOptions, error) {
-	proxyName := context.Args().First()
-	if proxyName == "" {
-		cli.ShowSubcommandHelp(context)
-		return nil, errorf("Proxy name is missing.\n")
-	}
-
-	toxicName := context.String("toxicName")
-
-	return &toxiproxy.ToxicOptions{
-		ProxyName: proxyName,
-		ToxicName: toxicName,
-	}, nil
-}
-
-func parseAddToxicParams(c *cli.Context) (*toxiproxy.ToxicOptions, error) {
-	result, err := parseToxicCommonParams(c)
-	if err != nil {
-		return nil, err
-	}
-
-	result.ToxicType, err = getArgOrFail(c, "type")
-	if err != nil {
-		return nil, err
-	}
-
-	upstream := c.Bool("upstream")
-	downstream := c.Bool("downstream")
-	if upstream && downstream {
-		return nil, errorf("Only one should be specified: upstream or downstream.\n")
-	}
-
-	stream := "downstream"
-	if upstream {
-		stream = "upstream"
-	}
-	result.Stream = stream
-
-	result.Toxicity, err = parseToxicity(c, 1.0)
-	if err != nil {
-		return nil, err
-	}
-
-	result.Attributes = parseAttributes(c, "attribute")
-
-	return result, nil
-}
-
 func addToxic(c *cli.Context, t *toxiproxy.Client) error {
 	toxicParams, err := parseAddToxicParams(c)
 	if err != nil {
@@ -518,60 +470,101 @@ func addToxic(c *cli.Context, t *toxiproxy.Client) error {
 }
 
 func updateToxic(c *cli.Context, t *toxiproxy.Client) error {
-	proxyName := c.Args().First()
-	if proxyName == "" {
-		cli.ShowSubcommandHelp(c)
-		return errorf("Proxy name is required as the first argument.\n")
-	}
-	toxicName, err := getArgOrFail(c, "toxicName")
+	toxicParams, err := parseUpdateToxicParams(c)
 	if err != nil {
 		return err
 	}
 
-	attributes := parseAttributes(c, "attribute")
-
-	p, err := t.Proxy(proxyName)
+	toxic, err := t.UpdateToxic(toxicParams)
 	if err != nil {
-		return errorf("Failed to retrieve proxy %s: %s\n", proxyName, err.Error())
+		return errorf("Failed to update toxic: %v\n", err)
 	}
 
-	toxicity, err := parseToxicity(c, -1)
+	fmt.Printf(
+		"Updated toxic '%s' on proxy '%s'\n",
+		toxic.Name,
+		toxicParams.ProxyName,
+	)
+	return nil
+}
+
+func removeToxic(c *cli.Context, t *toxiproxy.Client) error {
+	toxicParams, err := parseToxicCommonParams(c)
 	if err != nil {
-		return nil
+		return err
 	}
 
-	_, err = p.UpdateToxic(toxicName, toxicity, attributes)
+	err = t.RemoveToxic(toxicParams)
 	if err != nil {
-		return errorf("Failed to update toxic: %s\n", err.Error())
+		return errorf("Failed to remove toxic: %v\n", err)
 	}
 
-	fmt.Printf("Updated toxic '%s' on proxy '%s'\n", toxicName, proxyName)
+	fmt.Printf("Removed toxic '%s' on proxy '%s'\n", toxicParams.ToxicName, toxicParams.ProxyName)
 	return nil
 }
 
-func removeToxic(c *cli.Context, t *toxiproxy.Client) error {
-	proxyName := c.Args().First()
+func parseToxicCommonParams(context *cli.Context) (*toxiproxy.ToxicOptions, error) {
+	proxyName := context.Args().First()
 	if proxyName == "" {
-		cli.ShowSubcommandHelp(c)
-		return errorf("Proxy name is required as the first argument.\n")
+		cli.ShowSubcommandHelp(context)
+		return nil, errorf("Proxy name is missing.\n")
 	}
-	toxicName, err := getArgOrFail(c, "toxicName")
+
+	toxicName := context.String("toxicName")
+
+	return &toxiproxy.ToxicOptions{
+		ProxyName: proxyName,
+		ToxicName: toxicName,
+	}, nil
+}
+
+func parseUpdateToxicParams(c *cli.Context) (*toxiproxy.ToxicOptions, error) {
+	result, err := parseToxicCommonParams(c)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	p, err := t.Proxy(proxyName)
+	result.Toxicity, err = parseToxicity(c, 1.0)
 	if err != nil {
-		return errorf("Failed to retrieve proxy %s: %s\n", proxyName, err.Error())
+		return nil, err
 	}
 
-	err = p.RemoveToxic(toxicName)
+	result.Attributes = parseAttributes(c, "attribute")
+
+	return result, nil
+}
+
+func parseAddToxicParams(c *cli.Context) (*toxiproxy.ToxicOptions, error) {
+	result, err := parseToxicCommonParams(c)
 	if err != nil {
-		return errorf("Failed to remove toxic: %s\n", err.Error())
+		return nil, err
 	}
 
-	fmt.Printf("Removed toxic '%s' on proxy '%s'\n", toxicName, proxyName)
-	return nil
+	result.ToxicType, err = getArgOrFail(c, "type")
+	if err != nil {
+		return nil, err
+	}
+
+	upstream := c.Bool("upstream")
+	downstream := c.Bool("downstream")
+	if upstream && downstream {
+		return nil, errorf("Only one should be specified: upstream or downstream.\n")
+	}
+
+	stream := "downstream"
+	if upstream {
+		stream = "upstream"
+	}
+	result.Stream = stream
+
+	result.Toxicity, err = parseToxicity(c, 1.0)
+	if err != nil {
+		return nil, err
+	}
+
+	result.Attributes = parseAttributes(c, "attribute")
+
+	return result, nil
 }
 
 func parseAttributes(c *cli.Context, name string) toxiproxy.Attributes {

diff --git a/client/client.go b/client/client.go
@@ -189,6 +189,48 @@ func (client *Client) AddToxic(options *ToxicOptions) (*Toxic, error) {
 	return toxic, nil
 }
 
+// UpdateToxic update a toxic in proxy
+func (client *Client) UpdateToxic(options *ToxicOptions) (*Toxic, error) {
+	proxy, err := client.Proxy(options.ProxyName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve proxy with name `%s`: %v", options.ProxyName, err)
+	}
+
+	toxic, err := proxy.UpdateToxic(
+		options.ToxicName,
+		options.Toxicity,
+		options.Attributes,
+	)
+
+	if err != nil {
+		return nil,
+			fmt.Errorf(
+				"failed to update toxic '%s' of proxy '%s': %v",
+				options.ToxicName, options.ProxyName, err,
+			)
+	}
+
+	return toxic, nil
+}
+
+// RemoveToxic removes toxic from proxy
+func (client *Client) RemoveToxic(options *ToxicOptions) error {
+	proxy, err := client.Proxy(options.ProxyName)
+	if err != nil {
+		return fmt.Errorf("failed to retrieve proxy with name `%s`: %v", options.ProxyName, err)
+	}
+
+	err = proxy.RemoveToxic(options.ToxicName)
+	if err != nil {
+		return fmt.Errorf(
+			"failed to remove toxic '%s' from proxy '%s': %v",
+			options.ToxicName, options.ProxyName, err,
+		)
+	}
+
+	return nil
+}
+
 // Save saves changes to a proxy such as its enabled status or upstream port.
 func (proxy *Proxy) Save() error {
 	request, err := json.Marshal(proxy)