From f93741917503c3d4d0a6ffbcfaa0a34433c7e7e0 Mon Sep 17 00:00:00 2001
From: "Xinwei Xiong (cubxxw)" <3293172751nss@gmail.com>
Date: Thu, 28 Dec 2023 01:40:23 +0800
Subject: [PATCH] fix: fix openim zk env set

Signed-off-by: Xinwei Xiong (cubxxw) <3293172751nss@gmail.com>
---
 .../discoveryregister/discoveryregister_test.go    |  3 ++-
 .../discoveryregister/zookeeper/zookeeper.go       | 14 +++++++++++---
 tools/component/component_test.go                  | 12 ------------
 3 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/pkg/common/discoveryregister/discoveryregister_test.go b/pkg/common/discoveryregister/discoveryregister_test.go
index d4a634b91..d83da1285 100644
--- a/pkg/common/discoveryregister/discoveryregister_test.go
+++ b/pkg/common/discoveryregister/discoveryregister_test.go
@@ -24,7 +24,8 @@ import (
 
 func setupTestEnvironment() {
 	os.Setenv("ZOOKEEPER_SCHEMA", "openim")
-	os.Setenv("ZOOKEEPER_ADDRESS", "172.28.0.1:12181")
+	os.Setenv("ZOOKEEPER_ADDRESS", "172.28.0.1")
+	os.Setenv("ZOOKEEPER_PORT", "12181")
 	os.Setenv("ZOOKEEPER_USERNAME", "")
 	os.Setenv("ZOOKEEPER_PASSWORD", "")
 }
diff --git a/pkg/common/discoveryregister/zookeeper/zookeeper.go b/pkg/common/discoveryregister/zookeeper/zookeeper.go
index a4edffc43..db1a5c6c4 100644
--- a/pkg/common/discoveryregister/zookeeper/zookeeper.go
+++ b/pkg/common/discoveryregister/zookeeper/zookeeper.go
@@ -52,10 +52,18 @@ func getEnv(key, fallback string) string {
 	return fallback
 }
 
-// getZkAddrFromEnv returns the value of an environment variable if it exists, otherwise it returns the fallback value.
+// getZkAddrFromEnv returns the Zookeeper addresses combined from the ZOOKEEPER_ADDRESS and ZOOKEEPER_PORT environment variables.
+// If the environment variables are not set, it returns the fallback value.
 func getZkAddrFromEnv(fallback []string) []string {
-	if value, exists := os.LookupEnv("ZOOKEEPER_ADDRESS"); exists {
-		return strings.Split(value, ",")
+	address, addrExists := os.LookupEnv("ZOOKEEPER_ADDRESS")
+	port, portExists := os.LookupEnv("ZOOKEEPER_PORT")
+
+	if addrExists && portExists {
+		addresses := strings.Split(address, ",")
+		for i, addr := range addresses {
+			addresses[i] = addr + ":" + port
+		}
+		return addresses
 	}
 	return fallback
 }
diff --git a/tools/component/component_test.go b/tools/component/component_test.go
index 08ea0fda8..4488c029e 100644
--- a/tools/component/component_test.go
+++ b/tools/component/component_test.go
@@ -21,22 +21,10 @@ import (
 	"time"
 
 	"github.com/redis/go-redis/v9"
-	"github.com/stretchr/testify/assert"
 
 	"github.com/openimsdk/open-im-server/v3/pkg/common/config"
 )
 
-func TestCheckMysql(t *testing.T) {
-	err := mockInitCfg()
-	assert.NoError(t, err, "Initialization should not produce errors")
-
-	err = checkMysql()
-	if err != nil {
-		// You might expect an error if MySQL isn't running locally with the mock credentials.
-		t.Logf("Expected error due to mock configuration: %v", err)
-	}
-}
-
 // Mock for initCfg for testing purpose
 func mockInitCfg() error {
 	config.Config.Mysql.Username = "root"