feat: allow overriding registry rate limit in test command (argoproj-…

…labs#256) * feat: allow overriding registry rate limit in test command Signed-off-by: jannfis <[email protected]> * Add pflag to go.mod Signed-off-by: jannfis <[email protected]> * Fix linter issue Signed-off-by: jannfis <[email protected]> * go mod tidy Signed-off-by: jannfis <[email protected]>
sribiere-jellysmack · Aug 13, 2024 · 85ada96 · 85ada96
1 parent f7879c6
commit 85ada96
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/cmd/test.go b/cmd/test.go
@@ -9,6 +9,8 @@ import (
 	"github.com/argoproj-labs/argocd-image-updater/pkg/registry"
 
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"go.uber.org/ratelimit"
 )
 
 func newTestCommand() *cobra.Command {
@@ -23,6 +25,7 @@ func newTestCommand() *cobra.Command {
 		disableKubernetes bool
 		ignoreTags        []string
 		disableKubeEvents bool
+		rateLimit         int
 	)
 	var runCmd = &cobra.Command{
 		Use:   "test IMAGE",
@@ -102,6 +105,15 @@ argocd-image-updater test nginx --allow-tags '^1.19.\d+(\-.*)*$' --update-strate
 				log.Fatalf("could not set registry credentials: %v", err)
 			}
 
+			checkFlag := func(f *pflag.Flag) {
+				if f.Name == "rate-limit" {
+					log.Infof("Overriding registry rate-limit to %d requests per second", rateLimit)
+					ep.Limiter = ratelimit.New(rateLimit)
+				}
+			}
+
+			cmd.Flags().Visit(checkFlag)
+
 			var creds *image.Credential
 			var username, password string
 			if credentials != "" {
@@ -158,5 +170,6 @@ argocd-image-updater test nginx --allow-tags '^1.19.\d+(\-.*)*$' --update-strate
 	runCmd.Flags().StringVar(&kubeConfig, "kubeconfig", "", "path to your Kubernetes client configuration")
 	runCmd.Flags().StringVar(&credentials, "credentials", "", "the credentials definition for the test (overrides registry config)")
 	runCmd.Flags().BoolVar(&disableKubeEvents, "disable-kubernetes-events", false, "Disable kubernetes events")
+	runCmd.Flags().IntVar(&rateLimit, "rate-limit", 20, "specificy registry rate limit (overrides registry.conf)")
 	return runCmd
 }
diff --git a/go.mod b/go.mod
@@ -14,6 +14,7 @@ require (
 	github.com/prometheus/client_golang v1.7.1
 	github.com/sirupsen/logrus v1.7.0
 	github.com/spf13/cobra v1.1.3
+	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.7.0
 	go.uber.org/ratelimit v0.1.1-0.20201110185707-e86515f0dda9
 	golang.org/x/crypto v0.0.0-20210616213533-5ff15b29337e