Create helpers for iterating containers in a pod

pkg/api/v1/pod/util.go

@@ -48,6 +48,28 @@ func FindPort(pod *v1.Pod, svcPort *v1.ServicePort) (int, error) {
 	return 0, fmt.Errorf("no suitable port for manifest: %s", pod.UID)
 }
 
+// ContainerVisitor is called with each container spec, and returns true
+// if visiting should continue.
+type ContainerVisitor func(container *v1.Container) (shouldContinue bool)
+
+// VisitContainers invokes the visitor function with a pointer to the container
+// spec of every container in the given pod spec. If visitor returns false,
+// visiting is short-circuited. VisitContainers returns true if visiting completes,
+// false if visiting was short-circuited.
+func VisitContainers(podSpec *v1.PodSpec, visitor ContainerVisitor) bool {
+	for i := range podSpec.InitContainers {
+		if !visitor(&podSpec.InitContainers[i]) {
+			return false
+		}
+	}
+	for i := range podSpec.Containers {
+		if !visitor(&podSpec.Containers[i]) {
+			return false
+		}
+	}
+	return true
+}
+
 // Visitor is called with each object name, and returns true if visiting should continue
 type Visitor func(name string) (shouldContinue bool)
 
@@ -61,16 +83,9 @@ func VisitPodSecretNames(pod *v1.Pod, visitor Visitor) bool {
 			return false
 		}
 	}
-	for i := range pod.Spec.InitContainers {
-		if !visitContainerSecretNames(&pod.Spec.InitContainers[i], visitor) {
-			return false
-		}
-	}
-	for i := range pod.Spec.Containers {
-		if !visitContainerSecretNames(&pod.Spec.Containers[i], visitor) {
-			return false
-		}
-	}
+	VisitContainers(&pod.Spec, func(c *v1.Container) bool {
+		return visitContainerSecretNames(c, visitor)
+	})
 	var source *v1.VolumeSource
 
 	for i := range pod.Spec.Volumes {
@@ -152,16 +167,9 @@ func visitContainerSecretNames(container *v1.Container, visitor Visitor) bool {
 // Transitive references (e.g. pod -> pvc -> pv -> secret) are not visited.
 // Returns true if visiting completed, false if visiting was short-circuited.
 func VisitPodConfigmapNames(pod *v1.Pod, visitor Visitor) bool {
-	for i := range pod.Spec.InitContainers {
-		if !visitContainerConfigmapNames(&pod.Spec.InitContainers[i], visitor) {
-			return false
-		}
-	}
-	for i := range pod.Spec.Containers {
-		if !visitContainerConfigmapNames(&pod.Spec.Containers[i], visitor) {
-			return false
-		}
-	}
+	VisitContainers(&pod.Spec, func(c *v1.Container) bool {
+		return visitContainerConfigmapNames(c, visitor)
+	})
 	var source *v1.VolumeSource
 	for i := range pod.Spec.Volumes {
 		source = &pod.Spec.Volumes[i].VolumeSource

pkg/api/v1/pod/util_test.go

@@ -198,6 +198,92 @@ func TestFindPort(t *testing.T) {
 	}
 }
 
+func TestVisitContainers(t *testing.T) {
+	testCases := []struct {
+		description string
+		haveSpec    *v1.PodSpec
+		wantNames   []string
+	}{
+		{
+			"empty podspec",
+			&v1.PodSpec{},
+			[]string{},
+		},
+		{
+			"regular containers",
+			&v1.PodSpec{
+				Containers: []v1.Container{
+					{Name: "c1"},
+					{Name: "c2"},
+				},
+			},
+			[]string{"c1", "c2"},
+		},
+		{
+			"init containers",
+			&v1.PodSpec{
+				InitContainers: []v1.Container{
+					{Name: "i1"},
+					{Name: "i2"},
+				},
+			},
+			[]string{"i1", "i2"},
+		},
+		{
+			"regular and init containers",
+			&v1.PodSpec{
+				Containers: []v1.Container{
+					{Name: "c1"},
+					{Name: "c2"},
+				},
+				InitContainers: []v1.Container{
+					{Name: "i1"},
+					{Name: "i2"},
+				},
+			},
+			[]string{"i1", "i2", "c1", "c2"},
+		},
+		{
+			"dropping fields",
+			&v1.PodSpec{
+				Containers: []v1.Container{
+					{Name: "c1"},
+					{Name: "c2", SecurityContext: &v1.SecurityContext{}},
+				},
+				InitContainers: []v1.Container{
+					{Name: "i1"},
+					{Name: "i2", SecurityContext: &v1.SecurityContext{}},
+				},
+			},
+			[]string{"i1", "i2", "c1", "c2"},
+		},
+	}
+
+	for _, tc := range testCases {
+		gotNames := []string{}
+		VisitContainers(tc.haveSpec, func(c *v1.Container) bool {
+			gotNames = append(gotNames, c.Name)
+			if c.SecurityContext != nil {
+				c.SecurityContext = nil
+			}
+			return true
+		})
+		if !reflect.DeepEqual(gotNames, tc.wantNames) {
+			t.Errorf("VisitContainers() for test case %q visited containers %q, wanted to visit %q", tc.description, gotNames, tc.wantNames)
+		}
+		for _, c := range tc.haveSpec.Containers {
+			if c.SecurityContext != nil {
+				t.Errorf("VisitContainers() for test case %q: got SecurityContext %#v for container %v, wanted nil", tc.description, c.SecurityContext, c.Name)
+			}
+		}
+		for _, c := range tc.haveSpec.InitContainers {
+			if c.SecurityContext != nil {
+				t.Errorf("VisitContainers() for test case %q: got SecurityContext %#v for init container %v, wanted nil", tc.description, c.SecurityContext, c.Name)
+			}
+		}
+	}
+}
+
 func TestPodSecrets(t *testing.T) {
 	// Stub containing all possible secret references in a pod.
 	// The names of the referenced secrets match struct paths detected by reflection.