Merge pull request #2339 from Zhupku/mengzezhu/ut2

test: add UT for ListenEndpoint in utils.go

Author: k8s-ci-robot

pkg/csi-common/utils.go

@@ -41,24 +41,32 @@ func parseEndpoint(ep string) (string, string, error) {
 	}
 	return "", "", fmt.Errorf("Invalid endpoint: %v", ep)
 }
+
+var klogFatalf = func(format string, args ...interface{}) {
+	klog.Fatalf(format, args...)
+}
+
 func ListenEndpoint(endpoint string) (net.Listener, error) {
 	proto, addr, err := parseEndpoint(endpoint)
 	if err != nil {
-		klog.Fatal(err.Error())
+		klogFatalf("Invalid endpoint: %v", err)
+		return nil, err
 	}
 
 	if proto == "unix" {
 		if runtime.GOOS != "windows" {
 			addr = "/" + addr
 		}
 		if err := os.Remove(addr); err != nil && !os.IsNotExist(err) {
-			klog.Fatalf("Failed to remove %s, error: %s", addr, err.Error())
+			klogFatalf("Failed to remove %s, error: %s", addr, err.Error())
+			return nil, err
 		}
 	}
 
 	listener, err := net.Listen(proto, addr)
 	if err != nil {
-		klog.Fatalf("Failed to listen: %v", err)
+		klogFatalf("Failed to listen: %v", err)
+		return nil, err
 	}
 	return listener, err
 }

pkg/csi-common/utils_test.go

@@ -20,6 +20,8 @@ import (
 	"bytes"
 	"context"
 	"flag"
+	"os"
+	"runtime"
 	"testing"
 
 	"google.golang.org/grpc"
@@ -302,3 +304,65 @@ func TestGetLogLevel(t *testing.T) {
 		}
 	}
 }
+
+func TestListenEndpoint(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("Skip test on Windows")
+	}
+
+	originalKlogFatalf := klogFatalf
+	klogFatalf = func(_ string, _ ...interface{}) {}
+	defer func() { klogFatalf = originalKlogFatalf }()
+
+	tests := []struct {
+		name     string
+		endpoint string
+		filePath string
+		wantErr  bool
+	}{
+		{
+			name:     "unix socket",
+			endpoint: "unix:///tmp/csi.sock",
+			filePath: "/tmp/csi.sock",
+			wantErr:  false,
+		},
+		{
+			name:     "tcp socket",
+			endpoint: "tcp://127.0.0.1:0",
+			wantErr:  false,
+		},
+		{
+			name:     "invalid endpoint",
+			endpoint: "invalid://",
+			wantErr:  true,
+		},
+		{
+			name:     "invalid unix socket",
+			endpoint: "unix://does/not/exist",
+			wantErr:  true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			defer func() {
+				if r := recover(); r != nil {
+					if !tt.wantErr {
+						t.Errorf("ListenEndpoint() panicked unexpectedly: %v", r)
+					}
+				}
+			}()
+
+			got, err := ListenEndpoint(tt.endpoint)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Listen() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if err == nil {
+				got.Close()
+				if tt.filePath != "" {
+					os.Remove(tt.filePath)
+				}
+			}
+		})
+	}
+}