verify every layer exists - optimized

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

Author: coderbirju

integration/convert_test.go

@@ -101,6 +101,57 @@ 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) {
+		// Helper function to check if a digest has the required label
+		checkDigestLabels := func(digest string) error {
+			if labels, ok := digestMap[digest]; ok {
+				for _, label := range labels {
+					if strings.Contains(label, "containerd.io/gc.ref.content") {
+						return nil // Found the required label
+					}
+				}
+				return fmt.Errorf("digest %s does not have associated gc label", digest)
+			}
+			return fmt.Errorf("digest %s not found in digestMap", digest)
+		}
+
+		// Check all layer digests
+		for _, layer := range manifest.Layers {
+			if err := checkDigestLabels(layer.Digest.String()); err != nil {
+				return false, err
+			}
+		}
+
+		// Check config digest
+		if err := checkDigestLabels(manifest.Config.Digest.String()); err != nil {
+			return false, err
+		}
+
+		return true, nil
+	}
+
 	if originalDigest == convertedDigest {
 		t.Fatalf("conversion did not change the digest: %s", originalDigest)
 	}
@@ -199,13 +250,13 @@ 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())
+		digestLabelMap, err := getDigestLabelMapping(manifestDesc.Digest.String())
 		if err != nil {
 			t.Errorf("failed to get config digest from manifest: %v", err)
 			continue
 		}
 
-		_, err = verifyImageLayersContainsLabel(configDigest, manifest)
+		_, err = verifyImageLayersContainsLabel(digestLabelMap, manifest)
 		if err != nil {
 			t.Errorf("error verifying layers: %v", err)
 			continue
@@ -214,49 +265,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()