From 93d72c7c125bfa3196c9ba23854db8dcb78dad23 Mon Sep 17 00:00:00 2001
From: Nalin Dahyabhai <nalin@redhat.com>
Date: Tue, 23 Jan 2024 15:43:59 -0500
Subject: [PATCH] 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 <nalin@redhat.com>
---
 pkg/manifests/manifests.go      | 42 +++++++++++++++++++++++++++++++++
 pkg/manifests/manifests_test.go | 24 +++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/pkg/manifests/manifests.go b/pkg/manifests/manifests.go
index 8296faa82..775dca75a 100644
--- 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
index ae103944c..a95e83666 100644
--- 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)
+		},
+	)
+}