From a8d8c9497adfa972dc556d7862ed2a44005ad580 Mon Sep 17 00:00:00 2001
From: Hedayat Vatankhah <hedayat.fwd@gmail.com>
Date: Thu, 13 Jul 2023 18:15:02 +0330
Subject: [PATCH] Add more tests for liveness probes with default hostname &
 named ports

Signed-off-by: Hedayat Vatankhah <hedayat.fwd@gmail.com>
---
 pkg/specgen/generate/kube/play_test.go | 130 ++++++++++++++++++++++++-
 1 file changed, 129 insertions(+), 1 deletion(-)

diff --git a/pkg/specgen/generate/kube/play_test.go b/pkg/specgen/generate/kube/play_test.go
index 4099f9c1b4..d84d10b0fd 100644
--- a/pkg/specgen/generate/kube/play_test.go
+++ b/pkg/specgen/generate/kube/play_test.go
@@ -1247,6 +1247,131 @@ func TestHttpLivenessProbe(t *testing.T) {
 			true,
 			"http://localhost:80/",
 		},
+		{
+			"HttpLivenessProbeNamedPort",
+			specgen.SpecGenerator{},
+			v1.Container{
+				LivenessProbe: &v1.Probe{
+					Handler: v1.Handler{
+						HTTPGet: &v1.HTTPGetAction{
+							Port: intstr.FromString("httpPort"),
+						},
+					},
+				},
+				Ports: []v1.ContainerPort{
+					{Name: "servicePort", ContainerPort: 7000},
+					{Name: "httpPort", ContainerPort: 8000},
+				},
+			},
+			"always",
+			true,
+			"http://localhost:8000/",
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			err := setupLivenessProbe(&test.specGenerator, test.container, test.restartPolicy)
+			if err == nil {
+				assert.Equal(t, err == nil, test.succeed)
+				assert.Contains(t, test.specGenerator.ContainerHealthCheckConfig.HealthConfig.Test, test.expectedURL)
+			}
+		})
+	}
+}
+
+func TestTCPLivenessProbe(t *testing.T) {
+	tests := []struct {
+		name          string
+		specGenerator specgen.SpecGenerator
+		container     v1.Container
+		restartPolicy string
+		succeed       bool
+		expectedHost  string
+		expectedPort  string
+	}{
+		{
+			"TCPLivenessProbeNormal",
+			specgen.SpecGenerator{},
+			v1.Container{
+				LivenessProbe: &v1.Probe{
+					Handler: v1.Handler{
+						TCPSocket: &v1.TCPSocketAction{
+							Host: "127.0.0.1",
+							Port: intstr.FromInt(8080),
+						},
+					},
+				},
+			},
+			"always",
+			true,
+			"127.0.0.1",
+			"8080",
+		},
+		{
+			"TCPLivenessProbeHostUsesDefault",
+			specgen.SpecGenerator{},
+			v1.Container{
+				LivenessProbe: &v1.Probe{
+					Handler: v1.Handler{
+						TCPSocket: &v1.TCPSocketAction{
+							Port: intstr.FromInt(200),
+						},
+					},
+				},
+			},
+			"always",
+			true,
+			"localhost",
+			"200",
+		},
+		{
+			"TCPLivenessProbeUseNamedPort",
+			specgen.SpecGenerator{},
+			v1.Container{
+				LivenessProbe: &v1.Probe{
+					Handler: v1.Handler{
+						TCPSocket: &v1.TCPSocketAction{
+							Port: intstr.FromString("servicePort"),
+							Host: "myservice.domain.com",
+						},
+					},
+				},
+				Ports: []v1.ContainerPort{
+					{ContainerPort: 6000},
+					{Name: "servicePort", ContainerPort: 4000},
+					{Name: "2ndServicePort", ContainerPort: 3000},
+				},
+			},
+			"always",
+			true,
+			"myservice.domain.com",
+			"4000",
+		},
+		{
+			"TCPLivenessProbeInvalidPortName",
+			specgen.SpecGenerator{},
+			v1.Container{
+				LivenessProbe: &v1.Probe{
+					Handler: v1.Handler{
+						TCPSocket: &v1.TCPSocketAction{
+							Port: intstr.FromString("3rdservicePort"),
+							Host: "myservice.domain.com",
+						},
+					},
+				},
+				Ports: []v1.ContainerPort{
+					{ContainerPort: 6000},
+					{Name: "servicePort", ContainerPort: 4000},
+					{Name: "2ndServicePort", ContainerPort: 3000},
+				},
+			},
+			"always",
+			false,
+			"myservice.domain.com",
+			"4000",
+		},
 	}
 
 	for _, test := range tests {
@@ -1254,7 +1379,10 @@ func TestHttpLivenessProbe(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			err := setupLivenessProbe(&test.specGenerator, test.container, test.restartPolicy)
 			assert.Equal(t, err == nil, test.succeed)
-			assert.Contains(t, test.specGenerator.ContainerHealthCheckConfig.HealthConfig.Test, test.expectedURL)
+			if err == nil {
+				assert.Contains(t, test.specGenerator.ContainerHealthCheckConfig.HealthConfig.Test, test.expectedHost)
+				assert.Contains(t, test.specGenerator.ContainerHealthCheckConfig.HealthConfig.Test, test.expectedPort)
+			}
 		})
 	}
 }