From a2cbc11d59465d2bc5b4bab465b99279e6c1d569 Mon Sep 17 00:00:00 2001
From: Daneyon Hansen <daneyon.hansen@solo.io>
Date: Wed, 23 Apr 2025 11:22:53 -0700
Subject: [PATCH] EPP: Update GetRandomPod() to return nil if no pods exist

Signed-off-by: Daneyon Hansen <daneyon.hansen@solo.io>
---
 pkg/epp/handlers/request.go              |  3 ++
 pkg/epp/handlers/server.go               |  3 ++
 pkg/epp/handlers/streamingserver_test.go | 55 ++++++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/pkg/epp/handlers/request.go b/pkg/epp/handlers/request.go
index 9121b59af..8d30e543d 100644
--- a/pkg/epp/handlers/request.go
+++ b/pkg/epp/handlers/request.go
@@ -138,6 +138,9 @@ func (s *StreamingServer) HandleRequestHeaders(ctx context.Context, reqCtx *Requ
 		// The above PR will address endpoint admission, but currently any request without a body will be
 		// routed to a random upstream pod.
 		pod := GetRandomPod(s.datastore)
+		if pod == nil {
+			return errutil.Error{Code: errutil.Internal, Msg: "no pods available in datastore"}
+		}
 		pool, err := s.datastore.PoolGet()
 		if err != nil {
 			return err
diff --git a/pkg/epp/handlers/server.go b/pkg/epp/handlers/server.go
index 2e3a35fe7..5e23c7a0a 100644
--- a/pkg/epp/handlers/server.go
+++ b/pkg/epp/handlers/server.go
@@ -449,6 +449,9 @@ func RandomWeightedDraw(logger logr.Logger, model *v1alpha2.InferenceModel, seed
 
 func GetRandomPod(ds datastore.Datastore) *backendmetrics.Pod {
 	pods := ds.PodGetAll()
+	if len(pods) == 0 {
+		return nil
+	}
 	number := rand.Intn(len(pods))
 	pod := pods[number]
 	return pod.GetPod()
diff --git a/pkg/epp/handlers/streamingserver_test.go b/pkg/epp/handlers/streamingserver_test.go
index 72f7031a4..23d2b68fa 100644
--- a/pkg/epp/handlers/streamingserver_test.go
+++ b/pkg/epp/handlers/streamingserver_test.go
@@ -18,8 +18,14 @@ package handlers
 
 import (
 	"testing"
+	"time"
+
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
@@ -126,6 +132,55 @@ func TestRandomWeightedDraw(t *testing.T) {
 	}
 }
 
+func TestGetRandomPod(t *testing.T) {
+	tests := []struct {
+		name      string
+		storePods []*corev1.Pod
+		expectNil bool
+	}{
+		{
+			name:      "No pods available",
+			storePods: []*corev1.Pod{},
+			expectNil: true,
+		},
+		{
+			name: "Single pod available",
+			storePods: []*corev1.Pod{
+				{ObjectMeta: metav1.ObjectMeta{Name: "pod1"}},
+			},
+			expectNil: false,
+		},
+		{
+			name: "Multiple pods available",
+			storePods: []*corev1.Pod{
+				{ObjectMeta: metav1.ObjectMeta{Name: "pod1"}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "pod2"}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "pod3"}},
+			},
+			expectNil: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			pmf := metrics.NewPodMetricsFactory(&metrics.FakePodMetricsClient{}, time.Millisecond)
+			ds := datastore.NewDatastore(t.Context(), pmf)
+			for _, pod := range test.storePods {
+				ds.PodUpdateOrAddIfNotExist(pod)
+			}
+
+			gotPod := GetRandomPod(ds)
+
+			if test.expectNil && gotPod != nil {
+				t.Errorf("expected nil pod, got: %v", gotPod)
+			}
+			if !test.expectNil && gotPod == nil {
+				t.Errorf("expected non-nil pod, got nil")
+			}
+		})
+	}
+}
+
 func pointer(v int32) *int32 {
 	return &v
 }