fix(controller): agent image name parsing

Support image name with registry prefix containing port number or
without tag (default to latest)

Author: starbops

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
+}

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)
+			}
+		})
+	}
+}