test: add ut for re-write probe container

apecloud · Oct 23, 2023 · 6b42fb7 · 6b42fb7
1 parent 300b14f
commit 6b42fb7
Showing 1 changed file with 37 additions and 0 deletions.
diff --git 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)))
+		})
+
+	})
+
 })