From 41afbaeccaff78c007c309822497c3e2cea0527f Mon Sep 17 00:00:00 2001
From: Kislay Kishore <kislayk@google.com>
Date: Tue, 25 Jun 2024 13:12:36 +0530
Subject: [PATCH] Add constraints on number of args

---
 cmd/root.go      |  1 +
 cmd/root_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/cmd/root.go b/cmd/root.go
index 86a8b52588..33876497ea 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -40,6 +40,7 @@ func NewRootCmd(mountFn func(config cfg.Config) error) (*cobra.Command, error) {
 and access Cloud Storage buckets as local file systems. For a technical overview
 of Cloud Storage FUSE, see https://cloud.google.com/storage/docs/gcs-fuse.`,
 		Version: getVersion(),
+		Args:    cobra.RangeArgs(1, 2),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if cfgErr != nil {
 				return fmt.Errorf("error while parsing config: %w", cfgErr)
diff --git a/cmd/root_test.go b/cmd/root_test.go
index 4249050516..4fc17314f8 100644
--- a/cmd/root_test.go
+++ b/cmd/root_test.go
@@ -46,3 +46,50 @@ func TestValidConfig(t *testing.T) {
 
 	assert.Nil(t, cmd.Execute())
 }
+
+func TestTooManyCobraArgs(t *testing.T) {
+	tests := []struct {
+		name        string
+		args        []string
+		expectError bool
+	}{
+		{
+			name:        "Too many args",
+			args:        []string{"abc", "pqr", "xyz"},
+			expectError: true,
+		},
+		{
+			name:        "Too few args",
+			args:        []string{},
+			expectError: true,
+		},
+		{
+			name:        "One arg is okay",
+			args:        []string{"abc"},
+			expectError: false,
+		},
+		{
+			name:        "Two args is okay",
+			args:        []string{"abc", "pqr"},
+			expectError: false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			cmd, err := NewRootCmd(func(config cfg.Config) error { return nil })
+			if err != nil {
+				t.Fatalf("Error while creating the root command: %v", err)
+			}
+			cmd.SetArgs(tc.args)
+
+			err = cmd.Execute()
+
+			if tc.expectError {
+				assert.NotNil(t, err)
+			} else {
+				assert.Nil(t, err)
+			}
+		})
+	}
+}