From af57d89420127c07b9a2f13bbb1a99fc3eb466eb Mon Sep 17 00:00:00 2001
From: presbrey <presbrey@users.noreply.github.com>
Date: Thu, 30 Jan 2025 21:23:38 -0500
Subject: [PATCH] Extract and improve command-line argument parsing logic

---
 main.go      | 47 +++++++++++++++++----------
 main_test.go | 91 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+), 17 deletions(-)

diff --git a/main.go b/main.go
index 969a459..adb5a88 100644
--- a/main.go
+++ b/main.go
@@ -512,25 +512,38 @@ func (lm *LightyMux) Run(ctx context.Context, configFile string) error {
 	return nil
 }
 
-func main() {
-	var opts = new(Options)
+// parseConfig parses command line flags and environment variables into Options
+func parseArgs(args []string) (*Options, string, error) {
+	opts := new(Options)
 
-	if err := env.Parse(&opts); err != nil {
-		fmt.Printf("Error parsing env vars: %v\n", err)
-		os.Exit(1)
+	if err := env.Parse(opts); err != nil {
+		return nil, "", fmt.Errorf("error parsing env vars: %w", err)
+	}
+
+	// Create a new FlagSet for testing purposes
+	fs := flag.NewFlagSet("lightymux", flag.ContinueOnError)
+	fs.StringVar(&opts.HTTPAddr, "http", opts.HTTPAddr, "HTTP listen address (e.g., :8080)")
+	fs.BoolVar(&opts.Verbose, "verbose", opts.Verbose, "Enable verbose logging")
+	fs.BoolVar(&opts.LogRequests, "log-requests", opts.LogRequests, "Log incoming requests")
+	fs.BoolVar(&opts.LogResponses, "log-responses", opts.LogResponses, "Log outgoing responses")
+	fs.BoolVar(&opts.LogErrors, "log-errors", opts.LogErrors, "Log proxy errors")
+	fs.StringVar(&opts.LogFile, "log-file", opts.LogFile, "Log to file instead of stderr")
+
+	if err := fs.Parse(args); err != nil {
+		return nil, "", err
+	}
+
+	if fs.NArg() != 1 {
+		return nil, "", fmt.Errorf("exactly one config file argument is required")
 	}
 
-	// Allow command line flags to override env vars
-	flag.StringVar(&opts.HTTPAddr, "http", opts.HTTPAddr, "HTTP listen address (e.g., :8080)")
-	flag.BoolVar(&opts.Verbose, "verbose", opts.Verbose, "Enable verbose logging")
-	flag.BoolVar(&opts.LogRequests, "log-requests", opts.LogRequests, "Log incoming requests")
-	flag.BoolVar(&opts.LogResponses, "log-responses", opts.LogResponses, "Log outgoing responses")
-	flag.BoolVar(&opts.LogErrors, "log-errors", opts.LogErrors, "Log proxy errors")
-	flag.StringVar(&opts.LogFile, "log-file", opts.LogFile, "Log to file instead of stderr")
-	flag.Parse()
-
-	if flag.NArg() != 1 {
-		fmt.Println("Usage: ./lightymux [flags] <config_file>")
+	return opts, fs.Arg(0), nil
+}
+
+func main() {
+	opts, configFile, err := parseArgs(os.Args[1:])
+	if err != nil {
+		fmt.Printf("Error parsing configuration: %v\n", err)
 		flag.PrintDefaults()
 		os.Exit(1)
 	}
@@ -541,7 +554,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err := lm.Run(context.Background(), flag.Arg(0)); err != nil {
+	if err := lm.Run(context.Background(), configFile); err != nil {
 		fmt.Printf("Error running LightyMux: %v\n", err)
 		os.Exit(1)
 	}
diff --git a/main_test.go b/main_test.go
index f5bd914..e722a72 100644
--- a/main_test.go
+++ b/main_test.go
@@ -995,3 +995,94 @@ routes:
 		})
 	}
 }
+
+func TestParseArgs(t *testing.T) {
+	// Save original environment
+	origEnv := os.Environ()
+	defer func() {
+		os.Clearenv()
+		for _, e := range origEnv {
+			parts := strings.SplitN(e, "=", 2)
+			os.Setenv(parts[0], parts[1])
+		}
+	}()
+
+	tests := []struct {
+		name        string
+		args        []string
+		env         map[string]string
+		wantOpts    *Options
+		wantConfig  string
+		wantErrText string
+	}{
+		{
+			name: "valid args and config",
+			args: []string{"-http=:8080", "-verbose", "config.yaml"},
+			env:  map[string]string{},
+			wantOpts: &Options{
+				HTTPAddr:     ":8080",
+				Verbose:     true,
+				LogRequests: false,
+			},
+			wantConfig: "config.yaml",
+		},
+		{
+			name: "environment variables",
+			args: []string{"config.yaml"},
+			env: map[string]string{
+				"HTTP_ADDR":     ":9090",
+				"LOG_REQUESTS":  "true",
+				"LOG_RESPONSES": "true",
+			},
+			wantOpts: &Options{
+				HTTPAddr:      ":9090",
+				LogRequests:  true,
+				LogResponses: true,
+			},
+			wantConfig: "config.yaml",
+		},
+		{
+			name:        "missing config file",
+			args:        []string{"-http=:8080"},
+			env:         map[string]string{},
+			wantErrText: "exactly one config file argument is required",
+		},
+		{
+			name:        "too many arguments",
+			args:        []string{"config1.yaml", "config2.yaml"},
+			env:         map[string]string{},
+			wantErrText: "exactly one config file argument is required",
+		},
+		{
+			name:        "invalid flag",
+			args:        []string{"-invalid-flag", "config.yaml"},
+			env:         map[string]string{},
+			wantErrText: "flag provided but not defined: -invalid-flag",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Clear environment and set test values
+			os.Clearenv()
+			for k, v := range tt.env {
+				os.Setenv(k, v)
+			}
+
+			opts, config, err := parseArgs(tt.args)
+
+			if tt.wantErrText != "" {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tt.wantErrText)
+				return
+			}
+
+			assert.NoError(t, err)
+			assert.Equal(t, tt.wantConfig, config)
+			assert.Equal(t, tt.wantOpts.HTTPAddr, opts.HTTPAddr)
+			assert.Equal(t, tt.wantOpts.Verbose, opts.Verbose)
+			assert.Equal(t, tt.wantOpts.LogRequests, opts.LogRequests)
+			assert.Equal(t, tt.wantOpts.LogResponses, opts.LogResponses)
+		})
+	}
+}