refactor port forward

Author: Revolyssup

pkg/forwarder/forwarder.go

@@ -0,0 +1,5 @@
+package forwarder
+
+type Forwarder interface {
+	NextPort() string
+}

pkg/portforward/mock/mock.go

@@ -0,0 +1,22 @@
+package mock
+
+import "github.com/revolyssup/k8sdebug/pkg/forwarder"
+
+type MockForwarder struct {
+	ports []string
+	index int
+}
+
+// NewMock creates a mock forwarder with predefined ports.
+func New(ports ...string) forwarder.Forwarder {
+	return &MockForwarder{
+		ports: ports,
+	}
+}
+
+// Port returns the next predefined port or an error if exhausted.
+func (m *MockForwarder) NextPort() string {
+	port := m.ports[m.index]
+	m.index = (m.index + 1) % len(m.ports)
+	return port
+}

pkg/portforward/mock/mock_test.go

@@ -0,0 +1,19 @@
+package mock_test
+
+import (
+	"testing"
+
+	"github.com/revolyssup/k8sdebug/pkg/portforward/mock"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRoundRobin(t *testing.T) {
+	mock := mock.New("8080", "8081", "8082")
+
+	assert.Equal(t, "8080", mock.NextPort())
+	assert.Equal(t, "8081", mock.NextPort())
+	assert.Equal(t, "8082", mock.NextPort())
+	assert.Equal(t, "8080", mock.NextPort())
+	assert.Equal(t, "8081", mock.NextPort())
+	assert.Equal(t, "8082", mock.NextPort())
+}

pkg/portforward/portforward.go

@@ -14,6 +14,9 @@ import (
 	"time"
 
 	"github.com/revolyssup/k8sdebug/pkg"
+	"github.com/revolyssup/k8sdebug/pkg/forwarder"
+	"github.com/revolyssup/k8sdebug/pkg/portforward/mock"
+	"github.com/revolyssup/k8sdebug/pkg/portforward/roundrobin"
 	"github.com/spf13/cobra"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -45,52 +48,21 @@ func forwardToPod(hostConn net.Conn, podCon net.Conn) {
 	io.Copy(podCon, hostConn)
 }
 
-func getPodConnection(fw forwarder) (net.Conn, error) {
-	port := fw.Port()
+func getPodConnection(fw forwarder.Forwarder) (net.Conn, error) {
+	port := fw.NextPort()
 	podConn, err := net.Dial("tcp", fmt.Sprintf(":%s", port))
 	if err != nil {
 		return nil, err
 	}
 	return podConn, nil
 }
 
-type forwarder interface {
-	Port() string
-}
-
-type roundRobin struct {
-	connNumber int
-	mx         sync.Mutex
-}
-
-func (rr *roundRobin) Port() string {
-	rr.mx.Lock()
-	defer rr.mx.Unlock()
-
-	initial := rr.connNumber
-	for {
-		rr.connNumber = (rr.connNumber + 1) % len(connPool)
-		portNum := connPool[rr.connNumber]
-		if portNum != "" {
-			// Check if port is actually listening
-			conn, err := net.DialTimeout("tcp", ":"+portNum, 50*time.Millisecond)
-			if err == nil {
-				conn.Close()
-				fmt.Println("PORT RETURNED ", portNum)
-				return portNum
-			}
-		}
-		if rr.connNumber == initial {
-			break // Avoid infinite loop
-		}
-	}
-	return ""
-}
-
-func getForwarder(policy string) forwarder {
+func getForwarder(policy string) forwarder.Forwarder {
 	switch policy {
 	case "round-robin":
-		return &roundRobin{}
+		return roundrobin.New(connPool)
+	case "mock":
+		return mock.New()
 	}
 	return nil
 }

pkg/portforward/roundrobin/roundrobin.go

@@ -0,0 +1,45 @@
+package roundrobin
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/revolyssup/k8sdebug/pkg/forwarder"
+)
+
+type RoundRobin struct {
+	connNumber int
+	mx         sync.Mutex
+	connPool   []string
+}
+
+func New(connPool []string) forwarder.Forwarder {
+	return &RoundRobin{
+		connPool: connPool,
+	}
+}
+func (rr *RoundRobin) NextPort() string {
+	rr.mx.Lock()
+	defer rr.mx.Unlock()
+
+	initial := rr.connNumber
+	for {
+		rr.connNumber = (rr.connNumber + 1) % len(rr.connPool)
+		portNum := rr.connPool[rr.connNumber]
+		if portNum != "" {
+			// Check if port is actually listening
+			conn, err := net.DialTimeout("tcp", ":"+portNum, 50*time.Millisecond)
+			if err == nil {
+				conn.Close()
+				fmt.Println("PORT RETURNED ", portNum)
+				return portNum
+			}
+		}
+		if rr.connNumber == initial {
+			break // Avoid infinite loop
+		}
+	}
+	return ""
+}