diff --git a/cmd/controller/root.go b/cmd/controller/root.go index 9eff2db..6bead90 100644 --- a/cmd/controller/root.go +++ b/cmd/controller/root.go @@ -47,16 +47,10 @@ var rootCmd = &cobra.Command{ } }, Run: func(cmd *cobra.Command, args []string) { - var image *config.Image - imageTokens := strings.Split(agentImage, ":") - if len(imageTokens) == 2 { - image = config.NewImage(imageTokens[0], imageTokens[1]) - } else { - fmt.Fprintf(os.Stderr, "Error parse agent image name\n") - if err := cmd.Help(); err != nil { - os.Exit(1) - } - os.Exit(0) + image, err := parseImageNameAndTag(agentImage) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) } options := &config.ControllerOptions{ @@ -96,3 +90,26 @@ func init() { func execute() { cobra.CheckErr(rootCmd.Execute()) } + +func parseImageNameAndTag(image string) (*config.Image, error) { + idx := strings.LastIndex(image, ":") + + if idx == -1 { + return config.NewImage(image, "latest"), nil + } + + // If the last colon is immediately followed by the end of the string, it's invalid (no tag). + if idx == len(image)-1 { + return nil, fmt.Errorf("invalid image name: colon without tag") + } + + if strings.Count(image, ":") > 2 { + return nil, fmt.Errorf("invalid image name: multiple colons found") + } + + if idx <= strings.LastIndex(image, "/") { + return config.NewImage(image, "latest"), nil + } + + return config.NewImage(image[:idx], image[idx+1:]), nil +} diff --git a/cmd/controller/root_test.go b/cmd/controller/root_test.go new file mode 100644 index 0000000..6c37f0e --- /dev/null +++ b/cmd/controller/root_test.go @@ -0,0 +1,78 @@ +package main + +import ( + "testing" + + "github.com/harvester/vm-dhcp-controller/pkg/config" + "github.com/stretchr/testify/assert" +) + +func TestParseImageNameAndTag(t *testing.T) { + tests := []struct { + name string + image string + expected *config.Image + err bool + }{ + { + name: "valid image with registry and tag", + image: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller:v0.3.3", + expected: &config.Image{ + Repository: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller", + Tag: "v0.3.3", + }, + err: false, + }, + { + name: "valid image with only image and tag", + image: "rancher/harvester-vm-dhcp-controller:v0.3.3", + expected: &config.Image{ + Repository: "rancher/harvester-vm-dhcp-controller", + Tag: "v0.3.3", + }, + err: false, + }, + { + name: "valid image without tag", + image: "rancher/harvester-vm-dhcp-controller", + expected: &config.Image{ + Repository: "rancher/harvester-vm-dhcp-controller", + Tag: "latest", + }, + err: false, + }, + { + name: "valid image with port but no tag", + image: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller", + expected: &config.Image{ + Repository: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller", + Tag: "latest", + }, + err: false, + }, + { + name: "invalid image with colon but no tag", + image: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller:", + expected: nil, + err: true, + }, + { + name: "invalid image with multiple colons", + image: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller:v0.3.3:latest", + expected: nil, + err: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := parseImageNameAndTag(tt.image) + if tt.err { + assert.Error(t, err) + } else { + assert.NoError(t, err) + assert.Equal(t, tt.expected, result) + } + }) + } +}