add unit test in scheduler plugin part(by-label, data-parallel-profile-handler, pd-profile-handler) (#461)

Signed-off-by: CYJiang <googs1025@gmail.com>

Author: googs1025

pkg/plugins/filter/by_label_test.go

@@ -1,11 +1,17 @@
 package filter
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	k8stypes "k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
+	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 )
 
 func TestByLabelFactory(t *testing.T) {
@@ -130,3 +136,126 @@ func TestByLabelFactoryInvalidJSON(t *testing.T) {
 		})
 	}
 }
+
+// Helper functions
+func createPod(nsn k8stypes.NamespacedName, ipaddr string, labels map[string]string) types.Pod {
+	return &types.PodMetrics{
+		Pod: &backend.Pod{
+			NamespacedName: nsn,
+			Address:        ipaddr,
+			Labels:         labels,
+		},
+		MetricsState: &backendmetrics.MetricsState{},
+	}
+}
+
+func TestByLabelFiltering(t *testing.T) {
+	pods := []types.Pod{
+		createPod(k8stypes.NamespacedName{Namespace: "default", Name: "nginx-1"},
+			"10.0.0.1",
+			map[string]string{
+				"app":     "nginx",
+				"version": "v1.0",
+				"tier":    "frontend",
+			}),
+		createPod(k8stypes.NamespacedName{Namespace: "default", Name: "nginx-2"},
+			"10.0.0.2",
+			map[string]string{
+				"app":     "nginx",
+				"version": "v1.1",
+				"tier":    "frontend",
+			}),
+		createPod(k8stypes.NamespacedName{Namespace: "kube-system", Name: "coredns-1"},
+			"10.0.0.3",
+			map[string]string{
+				"app":  "coredns",
+				"tier": "system",
+			}),
+		createPod(k8stypes.NamespacedName{Namespace: "default", Name: "redis-1"},
+			"10.0.0.4",
+			map[string]string{
+				"app":        "redis",
+				"tier":       "backend",
+				"deprecated": "true",
+			}),
+		createPod(k8stypes.NamespacedName{Namespace: "default", Name: "web-1"},
+			"10.0.0.5",
+			map[string]string{
+				"app":         "web",
+				"tier":        "frontend",
+				"environment": "production",
+			}),
+		createPod(k8stypes.NamespacedName{Namespace: "default", Name: "no-tier-pod"},
+			"10.0.0.6",
+			map[string]string{
+				"app": "unknown",
+			}),
+	}
+
+	tests := []struct {
+		testName      string
+		label         string
+		validValues   []string
+		allowsNoLabel bool
+		expectedPods  []string // pod names that should match
+	}{
+		{
+			testName:      "match app nginx",
+			label:         "app",
+			validValues:   []string{"nginx"},
+			expectedPods:  []string{"nginx-1", "nginx-2"},
+			allowsNoLabel: false,
+		},
+		{
+			testName:      "match exact version v1.0",
+			label:         "version",
+			validValues:   []string{"v1.0"},
+			expectedPods:  []string{"nginx-1"},
+			allowsNoLabel: false,
+		},
+		{
+			testName:      "allow pods without 'tier' label",
+			label:         "tier",
+			allowsNoLabel: true,
+			expectedPods:  []string{"no-tier-pod"}, // only "no-tier-pod" doesn't have a "tier" label
+		},
+		{
+			testName:      "allow all known tier values",
+			label:         "tier",
+			validValues:   []string{"frontend", "backend", "system"},
+			allowsNoLabel: false,
+			expectedPods:  []string{"nginx-1", "nginx-2", "coredns-1", "redis-1", "web-1"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.testName, func(t *testing.T) {
+			rawParams, err := json.Marshal(byLabelParameters{
+				Label:         tt.label,
+				ValidValues:   tt.validValues,
+				AllowsNoLabel: tt.allowsNoLabel,
+			})
+			require.NoError(t, err)
+
+			plugin, err := ByLabelFactory("test-label", rawParams, nil)
+			require.NoError(t, err)
+			require.NotNil(t, plugin)
+
+			blf, ok := plugin.(*ByLabel)
+			require.True(t, ok, "plugin should be of type *ByLabel")
+
+			ctx := context.Background()
+			filteredPods := blf.Filter(ctx, nil, nil, pods)
+
+			var actualPodNames []string
+			for _, pod := range filteredPods {
+				actualPodNames = append(actualPodNames, pod.GetPod().NamespacedName.Name)
+			}
+
+			assert.ElementsMatch(t, tt.expectedPods, actualPodNames,
+				"filtered pods should match expected pods")
+			assert.Len(t, filteredPods, len(tt.expectedPods),
+				"filtered pods count should match expected count")
+		})
+	}
+}

pkg/plugins/profile/dp_profile_handler_test.go

@@ -1,10 +1,16 @@
 package profile
 
 import (
+	"context"
 	"encoding/json"
+	"net"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
+
+	"github.com/llm-d/llm-d-inference-scheduler/pkg/common"
 )
 
 func TestDataParallelProfileHandlerFactory(t *testing.T) {
@@ -113,3 +119,93 @@ func TestDataParallelProfileHandlerFactoryInvalidJSON(t *testing.T) {
 		})
 	}
 }
+
+func Test_DataParallelProfileHandler_ProcessResults(t *testing.T) {
+	tests := []struct {
+		name           string
+		primaryPort    int
+		profileResults map[string]*types.ProfileRunResult
+		expectError    bool
+		checkResult    func(*testing.T, *types.SchedulingResult, map[string]string)
+	}{
+		{
+			name: "error: multiple profiles not supported",
+			profileResults: map[string]*types.ProfileRunResult{
+				"profile1": newMockProfileRunResult(DefaultTestPodPort, "pod1"),
+				"profile2": newMockProfileRunResult(DefaultTestPodPort, "pod2"),
+			},
+			expectError: true,
+		},
+		{
+			name: "error: single profile but result is nil",
+			profileResults: map[string]*types.ProfileRunResult{
+				"nil-profile": nil,
+			},
+			expectError: true,
+		},
+		{
+			name:        "success: single profile with primaryPort → port overridden, header set",
+			primaryPort: 9000,
+			profileResults: map[string]*types.ProfileRunResult{
+				"dp-profile": newMockProfileRunResult(DefaultTestPodPort, "pod1"),
+			},
+			expectError: false,
+			checkResult: func(t *testing.T, res *types.SchedulingResult, headers map[string]string) {
+				assert.Equal(t, "dp-profile", res.PrimaryProfileName)
+
+				pods := res.ProfileResults["dp-profile"].TargetPods
+				require.Len(t, pods, 1)
+				assert.Equal(t, "9000", pods[0].GetPod().Port)                     // overridden
+				expectedHeader := net.JoinHostPort("10.0.0.1", DefaultTestPodPort) // original
+				assert.Equal(t, expectedHeader, headers[common.DataParallelPodHeader])
+			},
+		},
+		{
+			name:        "success: primaryPort=0 → port becomes '0'",
+			primaryPort: 0,
+			profileResults: map[string]*types.ProfileRunResult{
+				"dp": newMockProfileRunResult("8080", "pod1"),
+			},
+			expectError: false,
+			checkResult: func(t *testing.T, res *types.SchedulingResult, headers map[string]string) {
+				pod := res.ProfileResults["dp"].TargetPods[0]
+				assert.Equal(t, "0", pod.GetPod().Port)
+				assert.Equal(t, "10.0.0.1:8080", headers[common.DataParallelPodHeader])
+			},
+		},
+		{
+			name:        "success: multiple target pods → all ports overridden",
+			primaryPort: 8080,
+			profileResults: map[string]*types.ProfileRunResult{
+				"dp-profile": newMockProfileRunResult(DefaultTestPodPort, "pod1", "pod2"),
+			},
+			expectError: false,
+			checkResult: func(t *testing.T, res *types.SchedulingResult, headers map[string]string) {
+				pods := res.ProfileResults["dp-profile"].TargetPods
+				assert.Len(t, pods, 2)
+				for _, p := range pods {
+					assert.Equal(t, "8080", p.GetPod().Port)
+				}
+				assert.Equal(t, net.JoinHostPort("10.0.0.1", DefaultTestPodPort), headers[common.DataParallelPodHeader])
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			handler := NewDataParallelProfileHandler(tt.primaryPort).WithName("test-handler")
+			headers := make(map[string]string)
+			req := &types.LLMRequest{Headers: headers}
+			result, err := handler.ProcessResults(context.Background(), &types.CycleState{}, req, tt.profileResults)
+
+			if tt.expectError {
+				assert.Error(t, err)
+				return
+			}
+
+			require.NoError(t, err)
+			require.NotNil(t, result)
+			tt.checkResult(t, result, headers)
+		})
+	}
+}