From 6b42fb7ab055b093df8059ccf874d40bad157246 Mon Sep 17 00:00:00 2001
From: dereckchen <2564938992@qq.com>
Date: Mon, 23 Oct 2023 21:26:05 +0800
Subject: [PATCH] test: add ut for re-write probe container

---
 .../transformer_objection_generation_test.go  | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/pkg/controller/rsm/transformer_objection_generation_test.go b/pkg/controller/rsm/transformer_objection_generation_test.go
index 044bedf59e6..7a835878c75 100644
--- a/pkg/controller/rsm/transformer_objection_generation_test.go
+++ b/pkg/controller/rsm/transformer_objection_generation_test.go
@@ -178,4 +178,41 @@ var _ = Describe("object generation transformer test.", func() {
 			}
 		})
 	})
+
+	Context("injectRoleProbeBaseContainer function", func() {
+		It("should reuse container 'kb-checkrole' if exists", func() {
+			templateCopy := template.DeepCopy()
+			templateCopy.Spec.Containers = append(templateCopy.Spec.Containers, corev1.Container{
+				Name:  constant.RoleProbeContainerName,
+				Image: "bar",
+			})
+			injectRoleProbeBaseContainer(*rsm, templateCopy, "", nil)
+			Expect(len(templateCopy.Spec.Containers)).Should(Equal(2))
+			probeContainer := templateCopy.Spec.Containers[1]
+			Expect(probeContainer.ReadinessProbe).ShouldNot(BeNil())
+			Expect(len(probeContainer.Ports)).Should(Equal(1))
+			Expect(probeContainer.Ports[0].ContainerPort).Should(Equal(int32(defaultRoleProbeGRPCPort)))
+		})
+
+		It("should not use default grpcPort in case of 'lorry-grpc-port' existence", func() {
+			templateCopy := template.DeepCopy()
+			templateCopy.Spec.Containers = append(templateCopy.Spec.Containers, corev1.Container{
+				Name:  constant.RoleProbeContainerName,
+				Image: "bar",
+				Ports: []corev1.ContainerPort{
+					{
+						Name:          constant.LorryGRPCPortName,
+						ContainerPort: 9555,
+					},
+				},
+			})
+			injectRoleProbeBaseContainer(*rsm, templateCopy, "", nil)
+			Expect(len(templateCopy.Spec.Containers)).Should(Equal(2))
+			probeContainer := templateCopy.Spec.Containers[1]
+			Expect(len(probeContainer.Ports)).Should(Equal(1))
+			Expect(probeContainer.Ports[0].ContainerPort).Should(Equal(int32(9555)))
+		})
+
+	})
+
 })