diff --git a/internal/config/config.go b/internal/config/config.go index 2b6fdca..a3c3a6e 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -6,6 +6,7 @@ package config import ( "errors" "fmt" + "maps" "os" "path/filepath" "strings" @@ -25,6 +26,10 @@ var ( ErrInvalidTimeout = errors.New( "invalid config: timeout must be at least 1 when wait-for-ready is enabled, got", ) + ErrParamMissingWorkload = errors.New("missing workload prefix (expected workload.key=value)") + ErrParamEmptyWorkload = errors.New("empty workload name") + ErrParamMissingEquals = errors.New("missing '=' in param (expected workload.key=value)") + ErrParamEmptyKey = errors.New("empty param key") ) // WorkloadConfig holds per-workload configuration. @@ -130,6 +135,7 @@ func BindRunFlags(cmd *cobra.Command, defaultWorkloads []string) { f.StringSlice("ssh-key", nil, "SSH authorized key (repeatable)") f.StringSlice("ssh-key-file", nil, "SSH key file path (repeatable)") f.Int("vm-concurrency", constants.DefaultVMConcurrency, "Max concurrent VM creation operations") + f.String("params", "", "Per-workload params (comma-separated workload.key=value pairs)") } // BindCleanupFlags registers flags specific to the "cleanup" subcommand. @@ -232,6 +238,23 @@ func LoadConfig(cmd *cobra.Command) (*Config, error) { } cfg.Workloads = workloads + // Merge CLI/env params into workload configs (CLI flag > env var > YAML) + rawParams := resolveRawParams(cmd) + if rawParams != "" { + parsed, err := ParseParams(rawParams) + if err != nil { + return nil, fmt.Errorf("parsing --params: %w", err) + } + for wl, params := range parsed { + wlCfg := cfg.Workloads[wl] + if wlCfg.Params == nil { + wlCfg.Params = make(map[string]string) + } + maps.Copy(wlCfg.Params, params) + cfg.Workloads[wl] = wlCfg + } + } + if err := cfg.Validate(); err != nil { return nil, err } @@ -239,6 +262,45 @@ func LoadConfig(cmd *cobra.Command) (*Config, error) { return cfg, nil } +// ParseParams parses a comma-separated string of "workload.key=value" pairs +// into a per-workload param map. +func ParseParams(raw string) (map[string]map[string]string, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return make(map[string]map[string]string), nil + } + + result := make(map[string]map[string]string) + for pair := range strings.SplitSeq(raw, ",") { + pair = strings.TrimSpace(pair) + if pair == "" { + continue + } + + workload, rest, hasDot := strings.Cut(pair, ".") + if !hasDot { + return nil, fmt.Errorf("%w: %q", ErrParamMissingWorkload, pair) + } + if workload == "" { + return nil, fmt.Errorf("%w in %q", ErrParamEmptyWorkload, pair) + } + + key, value, hasEq := strings.Cut(rest, "=") + if !hasEq { + return nil, fmt.Errorf("%w: %q", ErrParamMissingEquals, pair) + } + if key == "" { + return nil, fmt.Errorf("%w in %q", ErrParamEmptyKey, pair) + } + + if result[workload] == nil { + result[workload] = make(map[string]string) + } + result[workload][key] = value + } + return result, nil +} + // bindFlagIfSet sets a Viper key from a Cobra flag only when the flag was explicitly provided. func bindFlagIfSet(v *viper.Viper, cmd *cobra.Command, name string) { if cmd.Flags().Changed(name) { @@ -297,3 +359,13 @@ func resolveSSHKeys(v *viper.Viper, cmd *cobra.Command) ([]string, error) { // Fall back to YAML config list return v.GetStringSlice("ssh-authorized-keys"), nil } + +// resolveRawParams returns the raw params string from CLI flag or env var. +// Priority: CLI --params flag > VIRTWORK_PARAMS env var. +func resolveRawParams(cmd *cobra.Command) string { + if cmd.Flags().Changed("params") { + val, _ := cmd.Flags().GetString("params") + return val + } + return os.Getenv("VIRTWORK_PARAMS") +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 1e3455c..e95510c 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -610,6 +610,7 @@ var _ = Describe("Config Validation", func() { }) }) + //nolint:dupl Context("timeout", func() { It("should reject zero timeout when wait-for-ready is enabled", func() { Expect(cmd.Flags().Set("timeout", "0")).To(Succeed()) @@ -641,3 +642,170 @@ var _ = Describe("Config Validation", func() { }) }) }) + +var _ = Describe("ParseParams", func() { + Context("happy paths", func() { + It("should parse a single workload.key=value pair", func() { + result, err := config.ParseParams("cpu.cpu-load-percent=50") + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(1)) + Expect(result["cpu"]).To(HaveKeyWithValue("cpu-load-percent", "50")) + }) + + It("should parse multiple pairs for the same workload", func() { + result, err := config.ParseParams("cpu.cpu-load-percent=50,cpu.cpu-method=matrixprod") + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(1)) + Expect(result["cpu"]).To(HaveKeyWithValue("cpu-load-percent", "50")) + Expect(result["cpu"]).To(HaveKeyWithValue("cpu-method", "matrixprod")) + }) + + It("should parse pairs across different workloads", func() { + result, err := config.ParseParams("cpu.cpu-load-percent=50,memory.vm-bytes=1G") + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(2)) + Expect(result["cpu"]).To(HaveKeyWithValue("cpu-load-percent", "50")) + Expect(result["memory"]).To(HaveKeyWithValue("vm-bytes", "1G")) + }) + + It("should return empty map for empty string", func() { + result, err := config.ParseParams("") + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(BeEmpty()) + }) + + It("should handle values containing equals signs", func() { + result, err := config.ParseParams("tps.extra-args=--flag=true") + Expect(err).NotTo(HaveOccurred()) + Expect(result["tps"]).To(HaveKeyWithValue("extra-args", "--flag=true")) + }) + + It("should trim whitespace around pairs", func() { + result, err := config.ParseParams(" cpu.cpu-load-percent=50 , cpu.cpu-method=matrixprod ") + Expect(err).NotTo(HaveOccurred()) + Expect(result["cpu"]).To(HaveKeyWithValue("cpu-load-percent", "50")) + Expect(result["cpu"]).To(HaveKeyWithValue("cpu-method", "matrixprod")) + }) + }) + + Context("error cases", func() { + It("should reject a pair with no dot separator", func() { + _, err := config.ParseParams("key=value") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("missing workload prefix")) + }) + + It("should reject a pair with empty workload name", func() { + _, err := config.ParseParams(".key=value") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("empty workload name")) + }) + + It("should reject a pair with no equals sign", func() { + _, err := config.ParseParams("cpu.key") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("missing '='")) + }) + + It("should reject a pair with empty key", func() { + _, err := config.ParseParams("cpu.=value") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("empty param key")) + }) + }) +}) + +var _ = Describe("--params flag integration", func() { + var cmd *cobra.Command + + BeforeEach(func() { + for _, env := range os.Environ() { + if strings.HasPrefix(env, "VIRTWORK_") { + _ = os.Unsetenv(strings.Split(env, "=")[0]) + } + } + cmd = newTestCommand() + }) + + It("should parse --params flag into workload configs", func() { + Expect(cmd.Flags().Set("params", "cpu.cpu-load-percent=50,cpu.cpu-method=matrixprod")).To(Succeed()) + + cfg, err := config.LoadConfig(cmd) + Expect(err).NotTo(HaveOccurred()) + Expect(cfg.Workloads).To(HaveKey("cpu")) + Expect(cfg.Workloads["cpu"].Params).To(HaveKeyWithValue("cpu-load-percent", "50")) + Expect(cfg.Workloads["cpu"].Params).To(HaveKeyWithValue("cpu-method", "matrixprod")) + }) + + It("should parse VIRTWORK_PARAMS env var", func() { + _ = os.Setenv("VIRTWORK_PARAMS", "memory.vm-bytes=2G") + defer func() { + _ = os.Unsetenv("VIRTWORK_PARAMS") + }() + + cfg, err := config.LoadConfig(cmd) + Expect(err).NotTo(HaveOccurred()) + Expect(cfg.Workloads).To(HaveKey("memory")) + Expect(cfg.Workloads["memory"].Params).To(HaveKeyWithValue("vm-bytes", "2G")) + }) + + It("should prefer --params flag over VIRTWORK_PARAMS env var", func() { + _ = os.Setenv("VIRTWORK_PARAMS", "cpu.cpu-load-percent=25") + defer func() { + _ = os.Unsetenv("VIRTWORK_PARAMS") + }() + + Expect(cmd.Flags().Set("params", "cpu.cpu-load-percent=75")).To(Succeed()) + + cfg, err := config.LoadConfig(cmd) + Expect(err).NotTo(HaveOccurred()) + Expect(cfg.Workloads["cpu"].Params).To(HaveKeyWithValue("cpu-load-percent", "75")) + }) + + Context("CLI params merge with YAML params", func() { + var tmpDir string + + BeforeEach(func() { + var err error + tmpDir, err = os.MkdirTemp("", "virtwork-config-test-*") + Expect(err).NotTo(HaveOccurred()) + }) + + AfterEach(func() { + _ = os.RemoveAll(tmpDir) + }) + + It("should override matching YAML keys and preserve others", func() { + path := writeConfigFile(tmpDir, ` +workloads: + cpu: + enabled: true + params: + cpu-load-percent: "80" + cpu-method: all +`) + Expect(cmd.Flags().Set("config", path)).To(Succeed()) + Expect(cmd.Flags().Set("params", "cpu.cpu-load-percent=50")).To(Succeed()) + + cfg, err := config.LoadConfig(cmd) + Expect(err).NotTo(HaveOccurred()) + Expect(cfg.Workloads["cpu"].Params).To(HaveKeyWithValue("cpu-load-percent", "50")) + Expect(cfg.Workloads["cpu"].Params).To(HaveKeyWithValue("cpu-method", "all")) + }) + + It("should add CLI params for workloads not in YAML", func() { + path := writeConfigFile(tmpDir, ` +workloads: + cpu: + enabled: true +`) + Expect(cmd.Flags().Set("config", path)).To(Succeed()) + Expect(cmd.Flags().Set("params", "memory.vm-bytes=1G")).To(Succeed()) + + cfg, err := config.LoadConfig(cmd) + Expect(err).NotTo(HaveOccurred()) + Expect(cfg.Workloads["memory"].Params).To(HaveKeyWithValue("vm-bytes", "1G")) + Expect(cfg.Workloads).To(HaveKey("cpu")) + }) + }) +})