From 201c00481cac88ff463a34569da582d7e0239575 Mon Sep 17 00:00:00 2001
From: Jennifer Power <barnabei.jennifer@gmail.com>
Date: Tue, 18 Oct 2022 18:56:35 -0400
Subject: [PATCH] refactor: adds helper functions in orasclient to remove
 unnecessary calculations

Signed-off-by: Jennifer Power <barnabei.jennifer@gmail.com>
---
 registryclient/orasclient/oras.go      | 54 +++++++++++++++++---------
 registryclient/orasclient/oras_test.go |  2 +-
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/registryclient/orasclient/oras.go b/registryclient/orasclient/oras.go
index a26da7f9..f6ec2630 100644
--- a/registryclient/orasclient/oras.go
+++ b/registryclient/orasclient/oras.go
@@ -120,21 +120,18 @@ func (c *orasClient) LoadCollection(ctx context.Context, reference string) (coll
 	if exists {
 		return value.(collection.Collection), nil
 	}
-
-	desc, _, err := c.GetManifest(ctx, reference)
+	repo, updatedRef, err := c.setupRepo(ctx, reference)
 	if err != nil {
-		return collection.Collection{}, err
-	}
-	fetcherFn := func(ctx context.Context, desc ocispec.Descriptor) ([]byte, error) {
-		return c.GetContent(ctx, reference, desc)
+		return collection.Collection{}, fmt.Errorf("could not create registry target: %w", err)
 	}
-	co := collection.New(reference)
-	if err := collectionloader.LoadFromManifest(ctx, co, fetcherFn, desc); err != nil {
+
+	co, err := loadCollection(ctx, repo, updatedRef)
+	if err != nil {
 		return collection.Collection{}, err
 	}
-	co.Location = reference
-	c.collections.Store(reference, *co)
-	return *co, nil
+	c.collections.Store(reference, co)
+
+	return co, nil
 }
 
 // Pull performs a copy of OCI artifacts to a local location from a remote location.
@@ -152,7 +149,7 @@ func (c *orasClient) Pull(ctx context.Context, reference string, store content.S
 		from = cache.New(repo, c.cache)
 	}
 
-	graph, err := c.LoadCollection(ctx, reference)
+	graph, err := loadCollection(ctx, repo, updatedRef)
 	if err != nil {
 		return ocispec.Descriptor{}, allDescs, err
 	}
@@ -242,7 +239,6 @@ func (c *orasClient) Push(ctx context.Context, store content.Store, reference st
 	if err != nil {
 		return ocispec.Descriptor{}, fmt.Errorf("could not create registry target: %w", err)
 	}
-
 	return oras.Copy(ctx, store, updatedRef, repo, updatedRef, c.copyOpts)
 }
 
@@ -261,11 +257,7 @@ func (c *orasClient) GetContent(ctx context.Context, reference string, desc ocis
 	if err != nil {
 		return nil, fmt.Errorf("could not create registry target: %w", err)
 	}
-	r, err := repo.Fetch(ctx, desc)
-	if err != nil {
-		return nil, err
-	}
-	return orascontent.ReadAll(r, desc)
+	return getContent(ctx, desc, repo)
 }
 
 // Store returns the source storage being used to store
@@ -417,3 +409,29 @@ func getDefaultMediaType(file string) (string, error) {
 	}
 	return mType.String(), nil
 }
+
+// loadCollection is a helper function that allows a collection to be loaded with a given repository.
+func loadCollection(ctx context.Context, repo registry.Repository, reference string) (collection.Collection, error) {
+	desc, _, err := repo.FetchReference(ctx, reference)
+	if err != nil {
+		return collection.Collection{}, err
+	}
+	fetcherFn := func(ctx context.Context, desc ocispec.Descriptor) ([]byte, error) {
+		return getContent(ctx, desc, repo)
+	}
+	co := collection.New(reference)
+	if err := collectionloader.LoadFromManifest(ctx, co, fetcherFn, desc); err != nil {
+		return collection.Collection{}, err
+	}
+	co.Location = reference
+	return *co, nil
+}
+
+// getContent is a helper function that allows content to be retrieved with a given repository.
+func getContent(ctx context.Context, desc ocispec.Descriptor, repo registry.Repository) ([]byte, error) {
+	r, err := repo.Fetch(ctx, desc)
+	if err != nil {
+		return nil, err
+	}
+	return orascontent.ReadAll(r, desc)
+}
diff --git a/registryclient/orasclient/oras_test.go b/registryclient/orasclient/oras_test.go
index 948fd39a..29f4c344 100644
--- a/registryclient/orasclient/oras_test.go
+++ b/registryclient/orasclient/oras_test.go
@@ -266,7 +266,7 @@ func TestPushPull(t *testing.T) {
 				{
 					Prefix: "test.server.com",
 					Endpoint: registryclient.Endpoint{
-						PlainHTTP: true,
+						PlainHTTP: false,
 						Location:  "test.server.com",
 					},
 					Mirrors: []registryclient.Endpoint{