From e7350d126eba5c3561157a5275020d3088792e38 Mon Sep 17 00:00:00 2001
From: llhuii <liulinghui@huawei.com>
Date: Wed, 29 Sep 2021 10:01:38 +0800
Subject: [PATCH] TopologyAwareHints: add getHintsByZone unit test

---
 .../endpointslice/topologycache/utils_test.go | 134 ++++++++++++++++++
 1 file changed, 134 insertions(+)

diff --git a/pkg/controller/endpointslice/topologycache/utils_test.go b/pkg/controller/endpointslice/topologycache/utils_test.go
index 1ecb911ba8d..8e911369ee9 100644
--- a/pkg/controller/endpointslice/topologycache/utils_test.go
+++ b/pkg/controller/endpointslice/topologycache/utils_test.go
@@ -193,3 +193,137 @@ func Test_getGivingAndReceivingZones(t *testing.T) {
 		})
 	}
 }
+
+func Test_getHintsByZone(t *testing.T) {
+	testCases := []struct {
+		name                 string
+		slice                discovery.EndpointSlice
+		allocatedHintsByZone EndpointZoneInfo
+		allocations          map[string]Allocation
+		expectedHintsByZone  map[string]int
+	}{{
+		name:                 "empty",
+		slice:                discovery.EndpointSlice{},
+		allocations:          map[string]Allocation{},
+		allocatedHintsByZone: EndpointZoneInfo{},
+		expectedHintsByZone:  map[string]int{},
+	}, {
+		name: "single zone hint",
+		slice: discovery.EndpointSlice{
+			Endpoints: []discovery.Endpoint{{
+				Zone:  utilpointer.StringPtr("zone-a"),
+				Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{{Name: "zone-a"}}},
+			}},
+		},
+		allocations: map[string]Allocation{
+			"zone-a": {Maximum: 3},
+		},
+		allocatedHintsByZone: EndpointZoneInfo{"zone-a": 1},
+		expectedHintsByZone: map[string]int{
+			"zone-a": 1,
+		},
+	}, {
+		name: "multiple zone hints",
+		slice: discovery.EndpointSlice{
+			Endpoints: []discovery.Endpoint{
+				{
+					Zone:  utilpointer.StringPtr("zone-a"),
+					Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{{Name: "zone-a"}}},
+				},
+				{
+					Zone:  utilpointer.StringPtr("zone-a"),
+					Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{{Name: "zone-b"}}},
+				},
+				{
+					Zone:  utilpointer.StringPtr("zone-b"),
+					Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{{Name: "zone-b"}}},
+				},
+			},
+		},
+		allocations: map[string]Allocation{
+			"zone-a": {Maximum: 3},
+			"zone-b": {Maximum: 3},
+			"zone-c": {Maximum: 3},
+		},
+		allocatedHintsByZone: EndpointZoneInfo{"zone-a": 1, "zone-b": 1, "zone-c": 1},
+		expectedHintsByZone: map[string]int{
+			"zone-a": 1,
+			"zone-b": 2,
+		},
+	}, {
+		name: "invalid by zones that no longer requires any allocations",
+		slice: discovery.EndpointSlice{
+			Endpoints: []discovery.Endpoint{
+				{
+					Zone:  utilpointer.StringPtr("zone-a"),
+					Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{{Name: "zone-non-existent"}}},
+				},
+			},
+		},
+		allocations: map[string]Allocation{
+			"zone-a": {Maximum: 3},
+		},
+		allocatedHintsByZone: EndpointZoneInfo{"zone-a": 1, "zone-b": 1, "zone-c": 1},
+		expectedHintsByZone:  nil,
+	}, {
+		name: "invalid by endpoints with nil hints",
+		slice: discovery.EndpointSlice{
+			Endpoints: []discovery.Endpoint{
+				{
+					Zone:  utilpointer.StringPtr("zone-a"),
+					Hints: nil,
+				},
+			},
+		},
+		allocations: map[string]Allocation{
+			"zone-a": {Maximum: 3},
+		},
+		allocatedHintsByZone: EndpointZoneInfo{},
+		expectedHintsByZone:  nil,
+	}, {
+		name: "invalid by endpoint with no hints",
+		slice: discovery.EndpointSlice{
+			Endpoints: []discovery.Endpoint{
+				{
+					Zone:  utilpointer.StringPtr("zone-a"),
+					Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{}},
+				},
+			},
+		},
+		allocations: map[string]Allocation{
+			"zone-a": {Maximum: 3},
+		},
+		allocatedHintsByZone: EndpointZoneInfo{},
+		expectedHintsByZone:  nil,
+	}, {
+		name: "invalid by hints that would make minimum allocations impossible",
+		slice: discovery.EndpointSlice{
+			Endpoints: []discovery.Endpoint{
+				{
+					Zone:  utilpointer.StringPtr("zone-a"),
+					Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{{Name: "zone-a"}}},
+				},
+				{
+					Zone:  utilpointer.StringPtr("zone-a"),
+					Hints: &discovery.EndpointHints{ForZones: []discovery.ForZone{{Name: "zone-a"}}},
+				},
+			},
+		},
+		allocations: map[string]Allocation{
+			"zone-a": {Maximum: 2},
+		},
+		allocatedHintsByZone: EndpointZoneInfo{"zone-a": 1},
+		expectedHintsByZone:  nil,
+	}}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			actualHintsByZone := getHintsByZone(&tc.slice, tc.allocatedHintsByZone, tc.allocations)
+
+			if !reflect.DeepEqual(actualHintsByZone, tc.expectedHintsByZone) {
+				// %#v for distinguishing between nil and empty map
+				t.Errorf("Expected %#v hints by zones, got %#v", tc.expectedHintsByZone, actualHintsByZone)
+			}
+		})
+	}
+}