jeffreyrubi
diff --git a/‎cmd/options.go‎
Lines changed: 29 additions & 0 deletions b/‎cmd/options.go‎
Lines changed: 29 additions & 0 deletions
diff --git a/‎cmd/options_test.go‎
Lines changed: 72 additions & 0 deletions b/‎cmd/options_test.go‎
Lines changed: 72 additions & 0 deletions
diff --git a/‎cmd/root.go‎
Lines changed: 5 additions & 1 deletion b/‎cmd/root.go‎
Lines changed: 5 additions & 1 deletion
@@ -0,0 +1,29 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cmd
+
+import (
+ "github.com/googleapis/genai-toolbox/internal/log"
+)
+
+// Option is a function that configures a Command.
+type Option func(*Command)
+
+// WithLogger overrides the default logger.
+func WithLogger(l log.Logger) Option {
+ return func(c *Command) {
+ c.logger = l
+	}
+}
@@ -0,0 +1,72 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cmd
+
+import (
+ "errors"
+ "io"
+ "testing"
+
+ "github.com/googleapis/genai-toolbox/internal/log"
+ "github.com/spf13/cobra"
+)
+
+func TestCommandOptions(t *testing.T) {
+ logger, err := log.NewStdLogger(io.Discard, io.Discard, "INFO")
+ if err != nil {
+ t.Errorf("fail to initialize logger: %v", err)
+	}
+ tcs := []struct {
+ desc string
+ isValid func(*Command) error
+ option Option
+	}{
+		{
+ desc: "with logger",
+ isValid: func(c *Command) error {
+ if c.logger != logger {
+ return errors.New("loggers do not match")
+				}
+ return nil
+			},
+ option: WithLogger(logger),
+		},
+	}
+ for _, tc := range tcs {
+ t.Run(tc.desc, func(t *testing.T) {
+ got, err := invokeProxyWithOption(tc.option)
+ if err != nil {
+ t.Fatal(err)
+			}
+ if err := tc.isValid(got); err != nil {
+ t.Errorf("option did not initialize command correctly: %v", err)
+			}
+		})
+	}
+}
+
+func invokeProxyWithOption(o Option) (*Command, error) {
+ c := NewCommand(o)
+ // Keep the test output quiet
+ c.SilenceUsage = true
+ c.SilenceErrors = true
+ // Disable execute behavior
+ c.RunE = func(*cobra.Command, []string) error {
+ return nil
+	}
+
+ err := c.Execute()
+ return c, err
+}
@@ -67,7 +67,7 @@ type Command struct {
 }
 
 // NewCommand returns a Command object representing an invocation of the CLI.
-func NewCommand() *Command {
+func NewCommand(opts ...Option) *Command {
  cmd := &Command{
  Command: &cobra.Command{
  Use:           "toolbox",
@@ -76,6 +76,10 @@ func NewCommand() *Command {
 		},
 	}
 
+ for _, o := range opts {
+ o(cmd)
+	}
+
  flags := cmd.Flags()
  flags.StringVarP(&cmd.cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.")
  flags.IntVarP(&cmd.cfg.Port, "port", "p", 5000, "Port the server will listen on.")