diff --git a/pkg/cmd/attestation/artifact/oci/oci.go b/pkg/cmd/attestation/artifact/oci/oci.go
index 7028b1c0c..064f3835c 100644
--- a/pkg/cmd/attestation/artifact/oci/oci.go
+++ b/pkg/cmd/attestation/artifact/oci/oci.go
@@ -6,7 +6,7 @@ import (
 
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/name"
-	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
 )
@@ -62,53 +62,3 @@ func NewLiveClient() *LiveClient {
 		Get:            remote.Get,
 	}
 }
-
-// func NewMockClient() *Client {
-// 	return &Client{
-// 		ParseReference: func(string, ...name.Option) (name.Reference, error) {
-// 			return name.Tag{}, nil
-// 		},
-// 		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
-// 			d := remote.Descriptor{}
-// 			d.Digest = v1.Hash{
-// 				Hex:       "1234567890abcdef",
-// 				Algorithm: "sha256",
-// 			}
-
-// 			return &d, nil
-// 		},
-// 	}
-// }
-
-// func NewReferenceFailClient() *Client {
-// 	return &Client{
-// 		ParseReference: func(string, ...name.Option) (name.Reference, error) {
-// 			return nil, fmt.Errorf("failed to parse reference")
-// 		},
-// 		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
-// 			return nil, nil
-// 		},
-// 	}
-// }
-
-// func NewAuthFailClient() *Client {
-// 	return &Client{
-// 		ParseReference: func(string, ...name.Option) (name.Reference, error) {
-// 			return name.Tag{}, nil
-// 		},
-// 		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
-// 			return nil, &transport.Error{Errors: []transport.Diagnostic{{Code: transport.UnauthorizedErrorCode}}}
-// 		},
-// 	}
-// }
-
-// func NewDeniedClient() *Client {
-// 	return &Client{
-// 		ParseReference: func(string, ...name.Option) (name.Reference, error) {
-// 			return name.Tag{}, nil
-// 		},
-// 		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
-// 			return nil, &transport.Error{Errors: []transport.Diagnostic{{Code: transport.DeniedErrorCode}}}
-// 		},
-// 	}
-// }
diff --git a/pkg/cmd/attestation/artifact/oci/oci_test.go b/pkg/cmd/attestation/artifact/oci/oci_test.go
new file mode 100644
index 000000000..48643a5b0
--- /dev/null
+++ b/pkg/cmd/attestation/artifact/oci/oci_test.go
@@ -0,0 +1,83 @@
+package oci
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
+	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetImageDigest_Success(t *testing.T) {
+	expectedDigest := v1.Hash{
+		Hex:       "1234567890abcdef",
+		Algorithm: "sha256",
+	}
+
+	c := LiveClient{
+		ParseReference: func(string, ...name.Option) (name.Reference, error) {
+			return name.Tag{}, nil
+		},
+		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
+			d := remote.Descriptor{}
+			d.Digest = expectedDigest
+
+			return &d, nil
+		},
+	}
+
+	digest, err := c.GetImageDigest("test")
+	require.NoError(t, err)
+	require.Equal(t, &expectedDigest, digest)
+}
+
+func TestGetImageDigest_ReferenceFail(t *testing.T) {
+	c := LiveClient{
+		ParseReference: func(string, ...name.Option) (name.Reference, error) {
+			return nil, fmt.Errorf("failed to parse reference")
+		},
+		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
+			return nil, nil
+		},
+	}
+
+	digest, err := c.GetImageDigest("test")
+	require.Error(t, err)
+	require.Nil(t, digest)
+}
+
+func TestGetImageDigest_AuthFail(t *testing.T) {
+	c := LiveClient{
+		ParseReference: func(string, ...name.Option) (name.Reference, error) {
+			return name.Tag{}, nil
+		},
+		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
+			return nil, &transport.Error{Errors: []transport.Diagnostic{{Code: transport.UnauthorizedErrorCode}}}
+		},
+	}
+
+	digest, err := c.GetImageDigest("test")
+	require.Error(t, err)
+	require.ErrorIs(t, err, ErrRegistryAuthz)
+	require.Nil(t, digest)
+}
+
+func TestGetImageDigest_Denied(t *testing.T) {
+	c := LiveClient{
+		ParseReference: func(string, ...name.Option) (name.Reference, error) {
+			return name.Tag{}, nil
+		},
+		Get: func(name.Reference, ...remote.Option) (*remote.Descriptor, error) {
+			return nil, &transport.Error{Errors: []transport.Diagnostic{{Code: transport.DeniedErrorCode}}}
+		},
+	}
+
+	digest, err := c.GetImageDigest("test")
+	require.Error(t, err)
+	require.ErrorIs(t, err, ErrDenied)
+	require.Nil(t, digest)
+}