Add new test cases for gateway server (#1217)

Signed-off-by: Modi Tamam <[email protected]>

Author: ModiCodeCraftsman

go.mod

@@ -23,7 +23,7 @@ require (
 	github.com/prometheus/common v0.55.0
 	github.com/ray-project/kuberay/ray-operator v1.2.1
 	github.com/redis/go-redis/v9 v9.6.1
-	github.com/stretchr/testify v1.9.0
+	github.com/stretchr/testify v1.10.0
 	google.golang.org/grpc v1.65.0
 	k8s.io/api v0.31.2
 	k8s.io/apiextensions-apiserver v0.31.2
@@ -87,6 +87,7 @@ require (
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/tidwall/gjson v1.14.4 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect

go.sum

@@ -174,10 +174,14 @@ github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 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/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
 github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
 github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=

pkg/plugins/gateway/algorithms/router.go

@@ -17,11 +17,14 @@ limitations under the License.
 package routingalgorithms
 
 import (
-	"github.com/vllm-project/aibrix/pkg/types"
+	"sync"
 
+	"github.com/vllm-project/aibrix/pkg/types"
 	"k8s.io/klog/v2"
 )
 
+var routerMu sync.RWMutex
+
 const (
 	RouterNotSet = ""
 )
@@ -31,6 +34,8 @@ var routerConstructor = map[types.RoutingAlgorithm]types.RouterProviderRegistrat
 
 // Validate validates if user provided routing routers is supported by gateway
 func Validate(algorithms string) (types.RoutingAlgorithm, bool) {
+	routerMu.RLock()
+	defer routerMu.RUnlock()
 	if _, ok := routerFactory[types.RoutingAlgorithm(algorithms)]; ok {
 		return types.RoutingAlgorithm(algorithms), ok
 	} else {
@@ -41,6 +46,8 @@ func Validate(algorithms string) (types.RoutingAlgorithm, bool) {
 // Select the user provided router provider supported by gateway, no error reported and fallback to random router
 // Call Validate before this function to ensure expected behavior.
 func Select(algorithms types.RoutingAlgorithm) types.RouterProviderFunc {
+	routerMu.RLock()
+	defer routerMu.RUnlock()
 	if provider, ok := routerFactory[algorithms]; ok {
 		return provider
 	} else {
@@ -50,6 +57,8 @@ func Select(algorithms types.RoutingAlgorithm) types.RouterProviderFunc {
 }
 
 func Register(algorithm types.RoutingAlgorithm, constructor types.RouterConstructor) {
+	routerMu.Lock()
+	defer routerMu.Unlock()
 	routerConstructor[algorithm] = func() types.RouterProviderFunc {
 		router, err := constructor()
 		if err != nil {
@@ -63,6 +72,8 @@ func Register(algorithm types.RoutingAlgorithm, constructor types.RouterConstruc
 }
 
 func Init() {
+	routerMu.Lock()
+	defer routerMu.Unlock()
 	for algorithm, constructor := range routerConstructor {
 		routerFactory[algorithm] = constructor()
 		klog.Infof("Registered router for %s", algorithm)

pkg/plugins/gateway/gateway_req_body.go

@@ -52,7 +52,7 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestID string, reques
 
 	// early reject if no pods are ready to accept request for a model
 	podsArr, err := s.cache.ListPodsByModel(model)
-	if err != nil || podsArr == nil || podsArr.Len() == 0 || utils.CountRoutablePods(podsArr.All()) == 0 {
+	if err != nil || podsArr == nil || utils.CountRoutablePods(podsArr.All()) == 0 {
 		klog.ErrorS(err, "no ready pod available", "requestID", requestID, "model", model)
 		return generateErrorResponse(envoyTypePb.StatusCode_ServiceUnavailable,
 			[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{