Allow mc-router to scale a backend StatefulSet while routing traffic to a proxy via proxyServerName (#512)

server/routes_test.go

@@ -3,8 +3,10 @@ package server
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func Test_routesImpl_FindBackendForServerAddress(t *testing.T) {
@@ -66,9 +68,9 @@ func Test_routesImpl_FindBackendForServerAddress(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			r := NewRoutes()
 
-			r.CreateMapping(tt.mapping.serverAddress, tt.mapping.backend, nil, nil, "")
+			r.CreateMapping(tt.mapping.serverAddress, tt.mapping.backend, "", nil, nil, "")
 
-			if got, server, _, _ := r.FindBackendForServerAddress(context.Background(), tt.args.serverAddress); got != tt.want {
+			if got, server, _, _, _ := r.FindBackendForServerAddress(context.Background(), tt.args.serverAddress); got != tt.want {
 				t.Errorf("routesImpl.FindBackendForServerAddress() = %v, want %v", got, tt.want)
 			} else {
 				assert.Equal(t, tt.mapping.serverAddress, server)
@@ -76,3 +78,68 @@ func Test_routesImpl_FindBackendForServerAddress(t *testing.T) {
 		})
 	}
 }
+
+func Test_routesImpl_ScaleKey(t *testing.T) {
+	DownScaler = NewDownScaler(context.Background(), false, 1*time.Second)
+
+	t.Run("scaleKey defaults to backend when empty", func(t *testing.T) {
+		r := NewRoutes()
+		r.CreateMapping("mc.example.com", "backend:25565", "", nil, nil, "")
+
+		_, _, scaleKey, _, _ := r.FindBackendForServerAddress(context.Background(), "mc.example.com")
+		assert.Equal(t, "backend:25565", scaleKey)
+	})
+
+	t.Run("scaleKey is set when provided", func(t *testing.T) {
+		r := NewRoutes()
+		r.CreateMapping("mc.example.com", "proxy:25577", "10.0.0.5:25565", nil, nil, "")
+
+		backend, _, scaleKey, _, _ := r.FindBackendForServerAddress(context.Background(), "mc.example.com")
+		assert.Equal(t, "proxy:25577", backend)
+		assert.Equal(t, "10.0.0.5:25565", scaleKey)
+	})
+
+	t.Run("GetSleepers matches on scaleKey not backend", func(t *testing.T) {
+		r := NewRoutes()
+		called := false
+		sleeper := func(ctx context.Context) error {
+			called = true
+			return nil
+		}
+
+		// Two routes with same proxy backend but different scaleKeys
+		r.CreateMapping("mc1.example.com", "proxy:25577", "10.0.0.1:25565", nil, sleeper, "")
+		r.CreateMapping("mc2.example.com", "proxy:25577", "10.0.0.2:25565", nil, nil, "")
+
+		sleepers := r.GetSleepers("10.0.0.1:25565")
+		require.Len(t, sleepers, 1)
+		_ = sleepers[0](context.Background())
+		assert.True(t, called)
+
+		// No sleeper for the second scaleKey since it has nil sleeper
+		sleepers = r.GetSleepers("10.0.0.2:25565")
+		assert.Empty(t, sleepers)
+
+		// No sleeper when querying by proxy backend address
+		sleepers = r.GetSleepers("proxy:25577")
+		assert.Empty(t, sleepers)
+	})
+
+	t.Run("default route scaleKey", func(t *testing.T) {
+		r := NewRoutes()
+		r.SetDefaultRoute("proxy:25577", "10.0.0.5:25565", nil, nil, "")
+
+		backend, scaleKey, _, _ := r.GetDefaultRoute()
+		assert.Equal(t, "proxy:25577", backend)
+		assert.Equal(t, "10.0.0.5:25565", scaleKey)
+	})
+
+	t.Run("default route scaleKey defaults to backend", func(t *testing.T) {
+		r := NewRoutes()
+		r.SetDefaultRoute("backend:25565", "", nil, nil, "")
+
+		backend, scaleKey, _, _ := r.GetDefaultRoute()
+		assert.Equal(t, "backend:25565", backend)
+		assert.Equal(t, "backend:25565", scaleKey)
+	})
+}