From 5fe2f77b1c6101b2c550f182ec4021d828ff5bd0 Mon Sep 17 00:00:00 2001
From: Robin Vanderstraeten <robin.vanderstraeten@klarrio.com>
Date: Fri, 26 Jul 2024 18:00:53 +0200
Subject: [PATCH] Apply suggestions from code review

---
 lifecycle.go      | 20 +++++++--------
 lifecycle_test.go | 62 ++++++++++++++++++-----------------------------
 2 files changed, 34 insertions(+), 48 deletions(-)

diff --git a/lifecycle.go b/lifecycle.go
index 1ce274b4ee..b61ca4fd78 100644
--- a/lifecycle.go
+++ b/lifecycle.go
@@ -201,20 +201,21 @@ var defaultLogConsumersHook = func(cfg *LogConsumerConfig) ContainerLifecycleHoo
 func checkPortsMapped(exposedAndMappedPorts nat.PortMap, exposedPorts []string) error {
 	portMap, _, err := nat.ParsePortSpecs(exposedPorts)
 	if err != nil {
-		return err
+		return fmt.Errorf("could not check mapped ports: %w", err)
 	}
 
-	for exposedPort, _ := range portMap {
+	for exposedPort := range portMap {
 		// having entries in exposedAndMappedPorts, where the key is the exposed port,
 		// and the value is the mapped port, means that the port has been already mapped.
 		if _, ok := exposedAndMappedPorts[exposedPort]; !ok {
 			// check if the port is mapped with the protocol (default is TCP)
-			if !strings.Contains(string(exposedPort), "/") {
-				exposedPort = nat.Port(fmt.Sprintf("%s/tcp", string(exposedPort)))
-				if _, ok := exposedAndMappedPorts[exposedPort]; !ok {
-					return fmt.Errorf("port %s is not mapped yet", exposedPort)
-				}
-			} else {
+			if strings.Contains(string(exposedPort), "/") {
+				return fmt.Errorf("port %s is not mapped yet", exposedPort)
+			}
+
+			// Port didn't have a type, default to tcp and retry.
+			exposedPort = nat.Port(fmt.Sprintf("%s/tcp", string(exposedPort)))
+			if _, ok := exposedAndMappedPorts[exposedPort]; !ok {
 				return fmt.Errorf("port %s is not mapped yet", exposedPort)
 			}
 		}
@@ -247,8 +248,7 @@ var defaultReadinessHook = func() ContainerLifecycleHooks {
 							return err
 						}
 
-						exposedAndMappedPorts := jsonRaw.NetworkSettings.Ports
-						return checkPortsMapped(exposedAndMappedPorts, dockerContainer.exposedPorts)
+						return checkPortsMapped(jsonRaw.NetworkSettings.Ports, dockerContainer.exposedPorts)
 					},
 					b,
 					func(err error, duration time.Duration) {
diff --git a/lifecycle_test.go b/lifecycle_test.go
index ba466f8e30..19cace94f4 100644
--- a/lifecycle_test.go
+++ b/lifecycle_test.go
@@ -457,78 +457,64 @@ func TestMergePortBindings(t *testing.T) {
 }
 
 func TestPortMappingCheck(t *testing.T) {
-	makePortMap := func(ports []string) nat.PortMap {
+	makePortMap := func(ports ...string) nat.PortMap {
 		out := make(nat.PortMap)
 		for _, port := range ports {
 			// We don't care about the actual binding in this test
-			out[nat.Port(port)] = make([]nat.PortBinding, 0)
+			out[nat.Port(port)] = nil
 		}
 		return out
 	}
 
-	tests := []struct {
-		name                  string
+	tests := map[string]struct {
 		exposedAndMappedPorts nat.PortMap
 		exposedPorts          []string
 		expectError           bool
 	}{
-		{
-			name:                  "No protocol given",
-			exposedAndMappedPorts: makePortMap([]string{"1024/tcp"}),
+		"no-protocol": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
 			exposedPorts:          []string{"1024"},
-			expectError:           false,
 		},
-		{
-			name:                  "Protocol given",
-			exposedAndMappedPorts: makePortMap([]string{"1024/tcp"}),
+		"protocol": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
 			exposedPorts:          []string{"1024/tcp"},
-			expectError:           false,
 		},
-		{
-			name:                  "Protocol and target port given",
-			exposedAndMappedPorts: makePortMap([]string{"1024/tcp"}),
+		"protocol-target-port": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
 			exposedPorts:          []string{"1024:1024/tcp"},
-			expectError:           false,
 		},
-		{
-			name:                  "Only target port given, no protocol",
-			exposedAndMappedPorts: makePortMap([]string{"1024/tcp"}),
+		"target-port": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
 			exposedPorts:          []string{"1024:1024"},
-			expectError:           false,
 		},
-		{
-			name:                  "Multiple ports with different variations",
-			exposedAndMappedPorts: makePortMap([]string{"1024/tcp", "1025/tcp", "1026/tcp"}),
+		"multiple-ports": {
+			exposedAndMappedPorts: makePortMap("1024/tcp", "1025/tcp", "1026/tcp"),
 			exposedPorts:          []string{"1024", "25:1025/tcp", "1026:1026"},
-			expectError:           false,
 		},
-		{
-			name:                  "No ports mapped yet",
-			exposedAndMappedPorts: makePortMap([]string{}),
+		"no-mapped-ports": {
+			exposedAndMappedPorts: makePortMap(),
 			exposedPorts:          []string{"1024"},
 			expectError:           true,
 		},
-		{
-			name:                  "The wrong port has been mapped",
-			exposedAndMappedPorts: makePortMap([]string{"1023/tcp"}),
+		"wrong-mapped-port": {
+			exposedAndMappedPorts: makePortMap("1023/tcp"),
 			exposedPorts:          []string{"1024"},
 			expectError:           true,
 		},
-		{
-			name:                  "A subset of ports have been mapped",
-			exposedAndMappedPorts: makePortMap([]string{"1024/tcp", "1025/tcp"}),
+		"subset-mapped-ports": {
+			exposedAndMappedPorts: makePortMap("1024/tcp", "1025/tcp"),
 			exposedPorts:          []string{"1024", "1025", "1026"},
 			expectError:           true,
 		},
 	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
 			err := checkPortsMapped(tt.exposedAndMappedPorts, tt.exposedPorts)
 			if tt.expectError {
-				assert.Error(t, err)
-			} else {
-				assert.NoError(t, err)
+				require.Error(t, err)
+				return
 			}
+			require.NoError(t, err)
 		})
 	}
 }