pkg/manifests.list: add SetMediaType(), SetArtifactType()

Add SetMediaType()/MediaType() and SetArtifactType()/ArtifactType() for modifying information about instances in the list. Signed-off-by: Nalin Dahyabhai <[email protected]>
kwilczynski · Mar 13, 2024 · 93d72c7 · 93d72c7
1 parent 3754808
commit 93d72c7
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 0 deletions.
diff --git a/pkg/manifests/manifests.go b/pkg/manifests/manifests.go
@@ -33,6 +33,10 @@ type List interface {
 	Features(instanceDigest digest.Digest) ([]string, error)
 	SetOSFeatures(instanceDigest digest.Digest, osFeatures []string) error
 	OSFeatures(instanceDigest digest.Digest) ([]string, error)
+	SetMediaType(instanceDigest digest.Digest, mediaType string) error
+	MediaType(instanceDigest digest.Digest) (string, error)
+	SetArtifactType(instanceDigest digest.Digest, artifactType string) error
+	ArtifactType(instanceDigest digest.Digest) (string, error)
 	Serialize(mimeType string) ([]byte, error)
 	Instances() []digest.Digest
 	OCIv1() *v1.Index
@@ -357,6 +361,44 @@ func (l *list) OSFeatures(instanceDigest digest.Digest) ([]string, error) {
 	return append([]string{}, oci.Platform.OSFeatures...), nil
 }
 
+// SetMediaType sets the MediaType field in the instance with the specified digest.
+func (l *list) SetMediaType(instanceDigest digest.Digest, mediaType string) error {
+	oci, err := l.findOCIv1(instanceDigest)
+	if err != nil {
+		return err
+	}
+	oci.MediaType = mediaType
+	return nil
+}
+
+// MediaType retrieves the MediaType field in the instance with the specified digest.
+func (l *list) MediaType(instanceDigest digest.Digest) (string, error) {
+	oci, err := l.findOCIv1(instanceDigest)
+	if err != nil {
+		return "", err
+	}
+	return oci.MediaType, nil
+}
+
+// SetArtifactType sets the ArtifactType field in the instance with the specified digest.
+func (l *list) SetArtifactType(instanceDigest digest.Digest, artifactType string) error {
+	oci, err := l.findOCIv1(instanceDigest)
+	if err != nil {
+		return err
+	}
+	oci.ArtifactType = artifactType
+	return nil
+}
+
+// ArtifactType retrieves the ArtifactType field in the instance with the specified digest.
+func (l *list) ArtifactType(instanceDigest digest.Digest) (string, error) {
+	oci, err := l.findOCIv1(instanceDigest)
+	if err != nil {
+		return "", err
+	}
+	return oci.ArtifactType, nil
+}
+
 // FromBlob builds a list from an encoded manifest list or image index.
 func FromBlob(manifestBytes []byte) (List, error) {
 	manifestType := manifest.GuessMIMEType(manifestBytes)

diff --git a/pkg/manifests/manifests_test.go b/pkg/manifests/manifests_test.go
@@ -364,3 +364,27 @@ func TestSerialize(t *testing.T) {
 		}
 	}
 }
+
+func TestMediaType(t *testing.T) {
+	testString(t,
+		[]string{v1.MediaTypeImageManifest, manifest.DockerV2Schema2MediaType},
+		func(l List, i digest.Digest, s string) error {
+			return l.SetMediaType(i, s)
+		},
+		func(l List, i digest.Digest) (string, error) {
+			return l.MediaType(i)
+		},
+	)
+}
+
+func TestArtifactType(t *testing.T) {
+	testString(t,
+		[]string{"text/plain", "application/octet-stream"},
+		func(l List, i digest.Digest, s string) error {
+			return l.SetArtifactType(i, s)
+		},
+		func(l List, i digest.Digest) (string, error) {
+			return l.ArtifactType(i)
+		},
+	)
+}