Repository: vllm-project/aibrix
File: pkg/plugins/gateway/algorithms/least_gpu_cache_test.go (Go)

Code before:
	"github.com/vllm-project/aibrix/pkg/cache"
	"github.com/vllm-project/aibrix/pkg/metrics"
	"github.com/vllm-project/aibrix/pkg/types"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestLeastGpuCache(t *testing.T) {
	tests := []struct {
		name            string
		readyPods       []*v1.Pod
		podModelMetrics map[string]map[string]metrics.MetricValue
		expectErr       bool
		expectMsgs      []string
	}{
		{
			name: "successful routing with least gpu cache",
			readyPods: []*v1.Pod{
				{ObjectMeta: metav1.ObjectMeta{Name: "p1"}, Status: v1.PodStatus{PodIP: "1.1.1.1",
					Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}}},
				{ObjectMeta: metav1.ObjectMeta{Name: "p2"}, Status: v1.PodStatus{PodIP: "2.2.2.2",
					Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}}},
				{ObjectMeta: metav1.ObjectMeta{Name: "p3"}, Status: v1.PodStatus{PodIP: "3.3.3.3",
					Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}}},
				{ObjectMeta: metav1.ObjectMeta{Name: "p4"}, Status: v1.PodStatus{PodIP: "4.4.4.4",
					Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}}},
			},
			podModelMetrics: map[string]map[string]metrics.MetricValue{
				"p1": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.1}},
				"p2": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.2}},
				"p3": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.3}},
				"p4": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.4}},
			},
			expectErr:  false,
			expectMsgs: []string{"1.1.1.1:8000"},
		},
		{
			name:            "no ready pods",
			readyPods:       []*v1.Pod{},
			podModelMetrics: map[string]map[string]metrics.MetricValue{},
			expectErr:       true,
		},
		{
			name: "multiple pods with same gpu cache",
			readyPods: []*v1.Pod{
				{ObjectMeta: metav1.ObjectMeta{Name: "p1"}, Status: v1.PodStatus{PodIP: "1.1.1.1",
					Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}}},
				{ObjectMeta: metav1.ObjectMeta{Name: "p2"}, Status: v1.PodStatus{PodIP: "2.2.2.2",
					Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}}},
				{ObjectMeta: metav1.ObjectMeta{Name: "p3"}, Status: v1.PodStatus{PodIP: "3.3.3.3",
					Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}}},

Review comment (suggestion):
This code repeats. Why not use this function instead: [https://github.com/vllm-project/aibrix/blob/main/pkg/plugins/gateway/algorithms/least_busy_time_test.go#L220](https://github.com/vllm-project/aibrix/blob/main/pkg/plugins/gateway/algorithms/least_busy_time_test.go?utm_source=chatgpt.com#L220)
less copy-paste.

Code after:
	"github.com/vllm-project/aibrix/pkg/cache"
	"github.com/vllm-project/aibrix/pkg/metrics"
	"github.com/vllm-project/aibrix/pkg/types"
	v1 "k8s.io/api/core/v1"
)

func TestLeastGpuCache(t *testing.T) {
	tests := []struct {
		name            string
		readyPods       []*v1.Pod
		podModelMetrics map[string]map[string]metrics.MetricValue
		expectErr       bool
		expectMsgs      []string
	}{
		{
			name: "successful routing with least gpu cache",
			readyPods: []*v1.Pod{
				newPod("p1", "1.1.1.1", true, map[string]string{
					"model.aibrix.ai/port": "8000",
				}),
				newPod("p2", "2.2.2.2", true, map[string]string{
					"model.aibrix.ai/port": "8000",
				}),
				newPod("p3", "3.3.3.3", true, map[string]string{
					"model.aibrix.ai/port": "8000",
				}),
				newPod("p4", "4.4.4.4", true, map[string]string{
					"model.aibrix.ai/port": "8000",
				}),
			},
			podModelMetrics: map[string]map[string]metrics.MetricValue{
				"p1": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.1}},
				"p2": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.2}},
				"p3": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.3}},
				"p4": {metrics.GPUCacheUsagePerc: &metrics.SimpleMetricValue{Value: 0.4}},
			},
			expectErr:  false,
			expectMsgs: []string{"1.1.1.1:8000"},
		},
		{
			name:            "no ready pods",
			readyPods:       []*v1.Pod{},
			podModelMetrics: map[string]map[string]metrics.MetricValue{},
			expectErr:       true,
		},
		{
			name: "multiple pods with same gpu cache",
			readyPods: []*v1.Pod{
				newPod("p1", "1.1.1.1", true, map[string]string{
					"model.aibrix.ai/port": "8000",
				}),