verify every layer exists - optimized

Signed-off-by: Arjun Raja Yogidas <arjunry@amazon.com>

Author: coderbirju

integration/convert_test.go

@@ -101,6 +101,50 @@ func TestConvertWithForceRecreateZtocs(t *testing.T) {
 func validateConversion(t *testing.T, sh *shell.Shell, originalDigest, convertedDigest string) {
 	t.Helper()
 
+	// Helper function to get digest label mapping
+	getDigestLabelMapping := func(manifestDigest string) (map[string][]string, error) {
+		digestMap := make(map[string][]string)
+
+		manifestLabelOutput := sh.O("ctr", "content", "label", manifestDigest)
+		manifestLabels := strings.Split(strings.TrimSpace(string(manifestLabelOutput)), ",")
+
+		if len(manifestLabels) <= 0 {
+			return nil, fmt.Errorf("manifest does not contain any labels")
+		}
+
+		for _, label := range manifestLabels {
+			keyVal := strings.Split(label, "=")
+			if len(keyVal) != 2 {
+				continue
+			}
+			digestMap[keyVal[1]] = append(digestMap[keyVal[1]], keyVal[0])
+		}
+		return digestMap, nil
+	}
+
+	// Helper function to verify image layers contain required labels
+	verifyImageLayersContainsLabel := func(digestMap map[string][]string, manifest ocispec.Manifest) (bool, error) {
+		for i := 0; i < len(manifest.Layers); i++ {
+			layerDigest := manifest.Layers[i].Digest.String()
+			if labels, ok := digestMap[layerDigest]; ok {
+				found := false
+				for _, label := range labels {
+					if strings.Contains(label, "containerd.io/gc.ref.content") {
+						found = true
+						break
+					}
+				}
+				if !found {
+					return false, fmt.Errorf("digest %s does not have associated gc label", layerDigest)
+				}
+			} else {
+				// Layer digest not found in digestMap
+				return false, fmt.Errorf("digest %s not found in digestMap", layerDigest)
+			}
+		}
+		return true, nil
+	}
+
 	if originalDigest == convertedDigest {
 		t.Fatalf("conversion did not change the digest: %s", originalDigest)
 	}
@@ -199,7 +243,7 @@ func validateConversion(t *testing.T, sh *shell.Shell, originalDigest, converted
 		}
 
 		// get labels associated with the manifest and extract the digest
-		configDigest, err := getDigestLabelMapping(sh, manifestDesc.Digest.String())
+		configDigest, err := getDigestLabelMapping(manifestDesc.Digest.String())
 		if err != nil {
 			t.Errorf("failed to get config digest from manifest: %v", err)
 			continue
@@ -214,49 +258,6 @@ func validateConversion(t *testing.T, sh *shell.Shell, originalDigest, converted
 
 }
 
-func getDigestLabelMapping(sh *shell.Shell, manifestDigest string) (map[string][]string, error) {
-	digestMap := make(map[string][]string)
-
-	manifestLabelOutput := sh.O("ctr", "content", "label", manifestDigest)
-	manifestLabels := strings.Split(strings.TrimSpace(string(manifestLabelOutput)), ",")
-
-	if len(manifestLabels) <= 0 {
-		return nil, fmt.Errorf("manifest does not contain any labels")
-	}
-
-	for _, label := range manifestLabels {
-		keyVal := strings.Split(label, "=")
-		if len(keyVal) != 2 {
-			continue
-		} else {
-			digestMap[keyVal[1]] = append(digestMap[keyVal[1]], keyVal[0])
-		}
-	}
-	return digestMap, nil
-}
-
-func verifyImageLayersContainsLabel(digestMap map[string][]string, manifest ocispec.Manifest) (bool, error) {
-	for i := 0; i < len(manifest.Layers); i = i + 1 {
-		layerDigest := manifest.Layers[i].Digest.String()
-		if labels, ok := digestMap[layerDigest]; ok {
-			found := false
-			for _, label := range labels {
-				if strings.Contains(label, "containerd.io/gc.ref.content") {
-					found = true
-					break
-				}
-			}
-			if !found {
-				return false, fmt.Errorf("digest %s does not have associated gc label", layerDigest)
-			}
-		} else {
-			// Layer digest not found in digestMap
-			return false, fmt.Errorf("digest %s not found in digestMap", layerDigest)
-		}
-	}
-	return true, nil
-}
-
 func TestConvert(t *testing.T) {
 	sh, done := newSnapshotterBaseShell(t)
 	defer done()