chore(tests): added new test and CI

.github/workflows/ci.yaml

@@ -0,0 +1,82 @@
+name: CI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    name: Build and Test
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v3
+
+    - name: Set up Go
+      uses: actions/setup-go@v4
+      with:
+        go-version: '1.23'
+        cache: true # Let the action handle caching automatically
+
+    - name: Run go mod tidy
+      run: |
+        go mod tidy
+        if [ -n "$(git diff go.mod go.sum)" ]; then
+          echo "go.mod or go.sum is not up to date. Please run 'go mod tidy' locally and commit the changes."
+          git diff go.mod go.sum
+          exit 1
+        fi
+
+    - name: Verify dependencies
+      run: go mod verify
+
+    - name: Build
+      run: go build -v ./...
+
+    - name: Run tests
+      run: go test -v ./...
+
+    - name: Run vet
+      run: go vet ./...
+
+  build-loadbalancer:
+    name: Build Loadbalancer Container
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v3
+
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v2
+
+    - name: Build Loadbalancer Container
+      run: make build-loadbalancer
+
+  build-sidecar:
+    name: Build Sidecar Container
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v3
+
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v2
+
+    - name: Build Sidecar Container
+      run: make build-sidecar
+
+  build-controller:
+    name: Build Controller Container
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v3
+
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v2
+
+    - name: Build Controller Container
+      run: make build-controller
+

cmd/loadbalancer/server/rebalance.go

@@ -13,7 +13,7 @@ func handleRebalanceLoop(router route.RouterImpl, connections map[string]*connec
 		select {
 		case hosts := <-router.RebalanceRequests():
 			slog.Debug("Received message to rebalance", "hosts", hosts)
-			upstreamHostsToConnectionTracker := make(map[string]*connection.Connection)
+			upstreamHostsToConnectionTracker := make(map[string]*connection.Connection, len(connections))
 			slog.Debug("Flat mapping ConnectionTracker to upstreamHosts", "connections", connections)
 			for _, connectionTracker := range connections {
 				upstreamHostsToConnectionTracker[connectionTracker.User()] = connectionTracker

cmd/loadbalancer/server/rebalance_test.go

@@ -0,0 +1,89 @@
+package server
+
+import (
+	"log/slog"
+	"lukas8219/websocket-operator/cmd/loadbalancer/connection"
+	"net"
+	"testing"
+	"time"
+)
+
+type MockRouter struct {
+	rebalanceChan chan [][2]string
+	*slog.Logger
+}
+
+func (m *MockRouter) RebalanceRequests() <-chan [][2]string {
+	return m.rebalanceChan
+}
+
+func (m *MockRouter) Route(string) string { return "" }
+func (m *MockRouter) Add([]string)        {}
+func (m *MockRouter) GetAllUpstreamHosts() []string {
+	return []string{}
+}
+func (m *MockRouter) InitializeHosts() error { return nil }
+
+type MockConnection struct {
+	*connection.Connection
+	upstreamHost       string
+	user               string
+	upstreamCancelChan chan struct{}
+}
+
+type MockProxier struct{}
+
+func (m *MockProxier) ProxyDownstreamToUpstream() (net.Conn, error) {
+	return nil, nil
+}
+
+func (m *MockProxier) ProxyUpstreamToDownstream() {
+
+}
+
+func NewMockConnection(user, upstreamHost string) *connection.Connection {
+	conn := &net.TCPConn{}
+	tracker := connection.NewTracker(user, upstreamHost, "downstream", conn)
+	return &connection.Connection{
+		Tracker: tracker,
+		Proxier: &MockProxier{},
+	}
+}
+
+func TestHandleRebalanceLoop(t *testing.T) {
+	mockRouter := &MockRouter{
+		rebalanceChan: make(chan [][2]string, 1),
+	}
+	connections := make(map[string]*connection.Connection)
+
+	go handleRebalanceLoop(mockRouter, connections)
+
+	t.Run("No connection tracker found", func(t *testing.T) {
+		mockRouter.rebalanceChan <- [][2]string{{"non-existent", "new-host:3000"}}
+	})
+
+	t.Run("No need to rebalance - same host", func(t *testing.T) {
+		mockConn := NewMockConnection("user1", "same-host:3000")
+		connections["user1"] = mockConn
+
+		mockRouter.rebalanceChan <- [][2]string{{"user1", "same-host:3000"}}
+	})
+
+	t.Run("Successfully received cancellation signal", func(t *testing.T) {
+		mockConn := NewMockConnection("user2", "old-host:3000")
+		connections["user2"] = mockConn
+
+		go func() {
+			mockConn.Tracker.UpstreamCancelChan() <- 1
+		}()
+		mockRouter.rebalanceChan <- [][2]string{{"user2", "new-host:3000"}}
+		time.Sleep(100 * time.Millisecond)
+
+		//Expect Switch Host + Handle to be called
+		if mockConn.UpstreamHost() != "new-host:3000" {
+			t.Errorf("Expected host to be updated to new-host:3000, got %s", mockConn.UpstreamHost())
+		}
+	})
+
+	//TODO: change interface to use io.ReadWriter and inject Dialer so we can test new connections
+}

go.mod

@@ -7,15 +7,13 @@ toolchain go1.23.7
 require (
 	github.com/buraksezer/consistent v0.10.0
 	github.com/gobwas/ws v1.4.0
-	github.com/tylertreat/BoomFilters v0.0.0-20210315201527-1a82519a3e43
 	github.com/zeebo/xxh3 v1.0.2
 	k8s.io/api v0.32.3
 	k8s.io/apimachinery v0.32.3
 	k8s.io/client-go v0.32.3
 )
 
 require (
-	github.com/d4l3k/messagediff v1.2.1 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/emicklei/go-restful/v3 v3.11.0 // indirect
 	github.com/fxamacker/cbor/v2 v2.7.0 // indirect

go.sum

@@ -1,8 +1,6 @@
 github.com/buraksezer/consistent v0.10.0 h1:hqBgz1PvNLC5rkWcEBVAL9dFMBWz6I0VgUCW25rrZlU=
 github.com/buraksezer/consistent v0.10.0/go.mod h1:6BrVajWq7wbKZlTOUPs/XVfR8c0maujuPowduSpZqmw=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
-github.com/d4l3k/messagediff v1.2.1 h1:ZcAIMYsUg0EAp9X+tt8/enBE/Q8Yd5kzPynLyKptt9U=
-github.com/d4l3k/messagediff v1.2.1/go.mod h1:Oozbb1TVXFac9FtSIxHBMnBCq2qeH/2KkEQxENCrlLo=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
@@ -91,8 +89,6 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/tylertreat/BoomFilters v0.0.0-20210315201527-1a82519a3e43 h1:QEePdg0ty2r0t1+qwfZmQ4OOl/MB2UXIeJSpIZv56lg=
-github.com/tylertreat/BoomFilters v0.0.0-20210315201527-1a82519a3e43/go.mod h1:OYRfF6eb5wY9VRFkXJH8FFBi3plw2v+giaIu7P054pM=
 github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
 github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=