From e170757a6b47095a29b14a28a1685928970aeb17 Mon Sep 17 00:00:00 2001
From: bizhiyuan <bizhiyuan@kylinos.cn>
Date: Wed, 11 Dec 2024 14:48:55 +0800
Subject: [PATCH] Add Keys function Optimize Values function Add the unitcase
 of Keys and Values functions

---
 models/heartbeat.go |  2 +-
 utils/utils.go      | 16 +++++++---
 utils/utils_test.go | 77 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 90 insertions(+), 5 deletions(-)
 create mode 100644 utils/utils_test.go

diff --git a/models/heartbeat.go b/models/heartbeat.go
index b5b90ba..cf6855d 100644
--- a/models/heartbeat.go
+++ b/models/heartbeat.go
@@ -375,7 +375,7 @@ func DeleteHeartbeat(hbInfo HBInfo) utils.GeneralResponse {
 	}
 
 	for _, hbInfo := range hbInfoList {
-		ip := utils.Values(hbInfo)
+		ip := utils.Values[string, string](hbInfo)
 		if status == 0 {
 			linkId, _ = GetRingIdFromIPOnline(ip[0])
 		} else {
diff --git a/utils/utils.go b/utils/utils.go
index 13fcfa2..a599b9e 100644
--- a/utils/utils.go
+++ b/utils/utils.go
@@ -156,14 +156,22 @@ func Contains(slice []string, str string) bool {
 	return false
 }
 
-func Values(m map[string]string) []string {
-	values := make([]string, 0, len(m))
-	for _, value := range m {
-		values = append(values, value)
+func Values[K comparable, V any](m map[K]V) []V {
+	values := make([]V, 0, len(m))
+	for _, v := range m {
+		values = append(values, v)
 	}
 	return values
 }
 
+func Keys[K comparable, V any](m map[K]V) []K {
+	keys := make([]K, 0, len(m))
+	for k := range m {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
 func RemoveByValue[T comparable](slice []T, value T) []T {
 	var newSlice []T
 	for _, v := range slice {
diff --git a/utils/utils_test.go b/utils/utils_test.go
new file mode 100644
index 0000000..3eefebe
--- /dev/null
+++ b/utils/utils_test.go
@@ -0,0 +1,77 @@
+package utils
+
+import (
+	"reflect"
+	"sort"
+	"testing"
+)
+
+func TestKeys(t *testing.T) {
+	m := map[string]int{
+		"key1": 1,
+		"key2": 2,
+		"key3": 1,
+	}
+	keys := Keys[string, int](m)
+	expectedKeys := []string{"key1", "key2", "key3"}
+	sort.Strings(keys)
+	sort.Strings(expectedKeys)
+	if !reflect.DeepEqual(keys, expectedKeys) {
+		t.Errorf("Keys() = %v, want %v", keys, expectedKeys)
+	}
+
+	// 创建一个map来记录期望值出现的次数
+	expectedCount := make(map[string]int)
+	for _, v := range expectedKeys {
+		expectedCount[v]++
+	}
+
+	// 检查返回的slice中的每个值是否都在期望的slice中
+	for _, v := range keys {
+		if count, ok := expectedCount[v]; !ok || count == 0 {
+			t.Errorf("Keys() contains unexpected value %v", v)
+		}
+		expectedCount[v]--
+	}
+
+	// 检查期望的slice中的每个值是否都被找到了
+	for _, count := range expectedCount {
+		if count != 0 {
+			t.Errorf("Keys() is missing value")
+		}
+	}
+
+}
+
+func TestValues(t *testing.T) {
+	testMap := map[string]int{
+		"key1": 1,
+		"key2": 2,
+		"key3": 1,
+	}
+
+	values := Values2[string, int](testMap)
+
+	expectedKeys := []int{1, 2, 1}
+
+	// 创建一个map来记录期望值出现的次数
+	expectedCount := make(map[int]int)
+	for _, v := range expectedKeys {
+		expectedCount[v]++
+	}
+
+	// 检查返回的slice中的每个值是否都在期望的slice中
+	for _, v := range values {
+		if count, ok := expectedCount[v]; !ok || count == 0 {
+			t.Errorf("Values() contains unexpected value %v", v)
+		}
+		expectedCount[v]--
+	}
+
+	// 检查期望的slice中的每个值是否都被找到了
+	for _, count := range expectedCount {
+		if count != 0 {
+			t.Errorf("Values() is missing value")
+		}
+	}
+}
-- 
Gitee