fix test

apecloud · Nov 14, 2024 · faedfaa · faedfaa
1 parent 842ebbf
commit faedfaa
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 9 deletions.
diff --git a/pkg/controller/component/kbagent_test.go b/pkg/controller/component/kbagent_test.go
@@ -127,8 +127,9 @@ var _ = Describe("kb-agent", func() {
 
 			c := kbAgentContainer()
 			Expect(c).ShouldNot(BeNil())
-			Expect(c.Ports).Should(HaveLen(1))
+			Expect(c.Ports).Should(HaveLen(2))
 			Expect(c.Ports[0].ContainerPort).Should(Equal(int32(kbagent.DefaultHTTPPort)))
+			Expect(c.Ports[1].ContainerPort).Should(Equal(int32(kbagent.DefaultStreamingPort)))
 		})
 
 		It("port - in use", func() {
@@ -142,8 +143,9 @@ var _ = Describe("kb-agent", func() {
 
 			c := kbAgentContainer()
 			Expect(c).ShouldNot(BeNil())
-			Expect(c.Ports).Should(HaveLen(1))
+			Expect(c.Ports).Should(HaveLen(2))
 			Expect(c.Ports[0].ContainerPort).Should(Equal(int32(kbagent.DefaultHTTPPort + 1)))
+			Expect(c.Ports[1].ContainerPort).Should(Equal(int32(kbagent.DefaultStreamingPort + 1)))
 		})
 
 		It("startup env", func() {
@@ -152,7 +154,7 @@ var _ = Describe("kb-agent", func() {
 
 			c := kbAgentContainer()
 			Expect(c).ShouldNot(BeNil())
-			Expect(c.Env).Should(HaveLen(6))
+			Expect(c.Env).Should(HaveLen(7)) // 4 + 3
 		})
 
 		It("action env", func() {
@@ -177,7 +179,7 @@ var _ = Describe("kb-agent", func() {
 
 			c := kbAgentContainer()
 			Expect(c).ShouldNot(BeNil())
-			Expect(c.Env).Should(HaveLen(8))
+			Expect(c.Env).Should(HaveLen(9)) // 2 + 4 + 3
 			Expect(reflect.DeepEqual(c.Env[0], env[0])).Should(BeTrue())
 			Expect(reflect.DeepEqual(c.Env[1], env[1])).Should(BeTrue())
 		})

diff --git a/pkg/controller/component/lifecycle/kbagent.go b/pkg/controller/component/lifecycle/kbagent.go
@@ -272,11 +272,14 @@ func (a *kbagent) callActionWithSelector(ctx context.Context, spec *appsv1.Actio
 	//  - timeout
 	var output []byte
 	for _, pod := range pods {
-		host, port, err := a.serverEndpoint(pod)
-		if err != nil {
-			return nil, errors.Wrapf(err, "pod %s is unavailable to execute action %s", pod.Name, lfa.name())
+		endpoint := func() (string, int32, error) {
+			host, port, err := a.serverEndpoint(pod)
+			if err != nil {
+				return "", 0, errors.Wrapf(err, "pod %s is unavailable to execute action %s", pod.Name, lfa.name())
+			}
+			return host, port, nil
 		}
-		cli, err := kbacli.NewClient(host, port)
+		cli, err := kbacli.NewClient(endpoint)
 		if err != nil {
 			return nil, err // mock client error
 		}

diff --git a/pkg/kbagent/client/client.go b/pkg/kbagent/client/client.go
@@ -54,11 +54,15 @@ func GetMockClient() Client {
 	return mockClient
 }
 
-func NewClient(host string, port int32) (Client, error) {
+func NewClient(endpoint func() (string, int32, error)) (Client, error) {
 	if mockClient != nil || mockClientError != nil {
 		return mockClient, mockClientError
 	}
 
+	host, port, err := endpoint()
+	if err != nil {
+		return nil, err
+	}
 	if host == "" && port == 0 {
 		return nil, nil
 	}