From f4667f8a36ececd68b98c18bf7ba8c1e8d4cd9fe Mon Sep 17 00:00:00 2001
From: Chen Lin <linch@vmware.com>
Date: Fri, 22 Mar 2024 16:49:54 +0800
Subject: [PATCH] Fix empty pod ip pool type check

---
 pkg/cloudprovider/vsphereparavirtual/cloud.go | 22 +-------
 .../vsphereparavirtual/config.go              | 23 ++++++++
 .../vsphereparavirtual/config_test.go         | 56 +++++++++++++++++++
 3 files changed, 82 insertions(+), 19 deletions(-)

diff --git a/pkg/cloudprovider/vsphereparavirtual/cloud.go b/pkg/cloudprovider/vsphereparavirtual/cloud.go
index 2eaf94e70..2c81002e0 100644
--- a/pkg/cloudprovider/vsphereparavirtual/cloud.go
+++ b/pkg/cloudprovider/vsphereparavirtual/cloud.go
@@ -46,12 +46,6 @@ const (
 
 	// CloudControllerManagerNS is the namespace for vsphere paravirtual cluster cloud provider
 	CloudControllerManagerNS = "vmware-system-cloud-provider"
-
-	// PublicIPPoolType allows Pod IP address routable outside of Tier 0 router.
-	PublicIPPoolType = "Public"
-
-	// PrivateIPPoolType allows Pod IP address routable within VPC router.
-	PrivateIPPoolType = "Private"
 )
 
 var (
@@ -111,19 +105,9 @@ func newVSphereParavirtual(cfg *cpcfg.Config) (*VSphereParavirtual, error) {
 func (cp *VSphereParavirtual) Initialize(clientBuilder cloudprovider.ControllerClientBuilder, stop <-chan struct{}) {
 	klog.V(0).Info("Initing vSphere Paravirtual Cloud Provider")
 
-	if vpcModeEnabled {
-		if podIPPoolType != PublicIPPoolType && podIPPoolType != PrivateIPPoolType {
-			klog.Fatalf("Pod IP Pool Type can be either Public or Private in VPC network, %s is not supported", podIPPoolType)
-		}
-
-		if podIPPoolType == "" {
-			podIPPoolType = PrivateIPPoolType
-		}
-	} else {
-		// NSX-T T1 or VDS network
-		if podIPPoolType != "" {
-			klog.Fatal("Pod IP Pool Type can be set only when the network is VPC")
-		}
+	err := checkPodIPPoolType(vpcModeEnabled, podIPPoolType)
+	if err != nil {
+		klog.Fatalf("Invalid IP pool type: %v", err)
 	}
 
 	ownerRef, err := readOwnerRef(VsphereParavirtualCloudProviderConfigPath)
diff --git a/pkg/cloudprovider/vsphereparavirtual/config.go b/pkg/cloudprovider/vsphereparavirtual/config.go
index 00f3b4650..76e6faf3c 100644
--- a/pkg/cloudprovider/vsphereparavirtual/config.go
+++ b/pkg/cloudprovider/vsphereparavirtual/config.go
@@ -47,6 +47,10 @@ const (
 	SupervisorServiceAccountNameEnv string = "SUPERVISOR_CLUSTER_SERVICEACCOUNT_SECRET_NAME"
 	// SupervisorAPIServerFQDN reads supervisor service API server's fully qualified domain name from env
 	SupervisorAPIServerFQDN string = "supervisor.default.svc"
+	// PublicIPPoolType allows Pod IP address routable outside of Tier 0 router.
+	PublicIPPoolType = "Public"
+	// PrivateIPPoolType allows Pod IP address routable within VPC router.
+	PrivateIPPoolType = "Private"
 )
 
 // SupervisorEndpoint is the supervisor cluster endpoint
@@ -135,3 +139,22 @@ func getRestConfig(svConfigPath string) (*rest.Config, error) {
 		BearerToken: string(token),
 	}, nil
 }
+
+func checkPodIPPoolType(vpcModeEnabled bool, podIPPoolType string) error {
+	if vpcModeEnabled {
+		if podIPPoolType == "" {
+			return errors.New("--pod-ip-pool-type is required in the NSX-T VPC network")
+		}
+
+		if podIPPoolType != PublicIPPoolType && podIPPoolType != PrivateIPPoolType {
+			return errors.New("--pod-ip-pool-type can be either Public or Private in NSX-T VPC network, " + podIPPoolType + " is not supported")
+
+		}
+	} else {
+		// NSX-T T1 or VDS network
+		if podIPPoolType != "" {
+			return errors.New("--pod-ip-pool-type can be set only when the network is VPC")
+		}
+	}
+	return nil
+}
diff --git a/pkg/cloudprovider/vsphereparavirtual/config_test.go b/pkg/cloudprovider/vsphereparavirtual/config_test.go
index fbd18215b..97d194f99 100644
--- a/pkg/cloudprovider/vsphereparavirtual/config_test.go
+++ b/pkg/cloudprovider/vsphereparavirtual/config_test.go
@@ -23,6 +23,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 )
@@ -267,6 +268,61 @@ func TestGetRestConfig(t *testing.T) {
 	}
 }
 
+func TestCheckPodIPPoolType(t *testing.T) {
+	tests := []struct {
+		vpcModeEnabled   bool
+		podIPPoolType    string
+		expectedErrorMsg string
+		name             string
+	}{
+		{
+			name:             "If VPC mode is not enabled, --pod-ip-pool-type should be empty",
+			vpcModeEnabled:   false,
+			podIPPoolType:    "",
+			expectedErrorMsg: "",
+		},
+		{
+			name:             "If VPC mode is not enabled, throw out error if --pod-ip-pool-type is not empty",
+			vpcModeEnabled:   false,
+			podIPPoolType:    "test-ns",
+			expectedErrorMsg: "--pod-ip-pool-type can be set only when the network is VPC",
+		},
+		{
+			name:             "If VPC mode is enabled, throw error if --pod-ip-pool-type is not Public or Private",
+			vpcModeEnabled:   true,
+			podIPPoolType:    "test-ns",
+			expectedErrorMsg: "--pod-ip-pool-type can be either Public or Private in NSX-T VPC network, test-ns is not supported",
+		},
+		{
+			name:             "If VPC mode is enabled, throw error if --pod-ip-pool-type is empty",
+			vpcModeEnabled:   true,
+			podIPPoolType:    "",
+			expectedErrorMsg: "--pod-ip-pool-type is required in the NSX-T VPC network",
+		},
+		{
+			name:             "Pod IP Pool type should be successfully set as Public",
+			vpcModeEnabled:   true,
+			podIPPoolType:    "Public",
+			expectedErrorMsg: "",
+		},
+		{
+			name:             "Pod IP Pool type should be successfully set as Private",
+			vpcModeEnabled:   true,
+			podIPPoolType:    "Private",
+			expectedErrorMsg: "",
+		},
+	}
+
+	for _, test := range tests {
+		err := checkPodIPPoolType(test.vpcModeEnabled, test.podIPPoolType)
+		if test.expectedErrorMsg == "" {
+			assert.Equal(t, err, nil)
+		} else {
+			assert.Equal(t, err.Error(), test.expectedErrorMsg)
+		}
+	}
+}
+
 func createTestFile(dir, filename, content string) error {
 	tmpFile, err := os.Create(dir + "/" + filename)
 	if err != nil {