From b36a3325d0e22b3f74f45554c1dab654ed1e9fde Mon Sep 17 00:00:00 2001
From: Sergey-Kizimov <sergey.kizimov@hiya.com>
Date: Tue, 20 Aug 2024 19:57:34 -0700
Subject: [PATCH] Fill in the Size field in OCI downloaderResponse

Signed-off-by: Sergey-Kizimov <sergey.kizimov@hiya.com>
---
 download/oci_download.go      |  8 ++++++--
 download/oci_download_test.go | 24 ++++++++++++++++++------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/download/oci_download.go b/download/oci_download.go
index 89ae73d799c..b514a41b58b 100644
--- a/download/oci_download.go
+++ b/download/oci_download.go
@@ -211,7 +211,7 @@ func (d *OCIDownloader) oneShot(ctx context.Context) error {
 	d.SetCache(resp.etag) // set the current etag sha to the cache
 
 	if d.f != nil {
-		d.f(ctx, Update{ETag: resp.etag, Bundle: resp.b, Error: nil, Metrics: m, Raw: resp.raw})
+		d.f(ctx, Update{ETag: resp.etag, Bundle: resp.b, Error: nil, Metrics: m, Raw: resp.raw, Size: resp.size})
 	}
 	return nil
 }
@@ -258,7 +258,10 @@ func (d *OCIDownloader) download(ctx context.Context, m metrics.Metrics) (*downl
 		}, nil
 	}
 	fileReader, err := os.Open(bundleFilePath)
-	tee := io.TeeReader(fileReader, &buf)
+
+	cnt := &count{}
+	r := io.TeeReader(fileReader, cnt)
+	tee := io.TeeReader(r, &buf)
 
 	if err != nil {
 		return nil, err
@@ -281,6 +284,7 @@ func (d *OCIDownloader) download(ctx context.Context, m metrics.Metrics) (*downl
 		raw:      &buf,
 		etag:     etag,
 		longPoll: false,
+		size:     cnt.Bytes(),
 	}, nil
 }
 
diff --git a/download/oci_download_test.go b/download/oci_download_test.go
index 3b37935b937..86b68c1cc29 100644
--- a/download/oci_download_test.go
+++ b/download/oci_download_test.go
@@ -4,11 +4,11 @@
 package download
 
 import (
-	"bytes"
 	"context"
 	"encoding/base64"
 	"fmt"
 	"net/http"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -383,8 +383,7 @@ func TestOCICustomAuthPlugin(t *testing.T) {
 	}
 }
 
-func TestOCIRawDataHasNonZeroSize(t *testing.T) {
-	var buf bytes.Buffer
+func TestOCIValidateAndInjectDefaults(t *testing.T) {
 	ctx := context.Background()
 	fixture := newTestFixture(t)
 	fixture.server.expEtag = "sha256:c5834dbce332cabe6ae68a364de171a50bf5b08024c27d7c08cc72878b4df7ff"
@@ -407,14 +406,27 @@ func TestOCIRawDataHasNonZeroSize(t *testing.T) {
 
 	u1 := <-updates
 
-	buf.ReadFrom(u1.Raw)
+	if u1.Size == 0 {
+		t.Fatal("expected non-0 size")
+	}
 
 	if u1.Raw == nil {
 		t.Fatal("expected bundle reader to be non-nil")
 	}
 
-	if buf.Len() == 0 {
-		t.Fatal("expected non-0 size")
+	r := bundle.NewReader(u1.Raw)
+
+	b, err := r.Read()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !reflect.DeepEqual(b.Data, u1.Bundle.Data) {
+		t.Fatal("expected the bundle object and reader to have the same data")
+	}
+
+	if len(b.Modules) != len(u1.Bundle.Modules) {
+		t.Fatal("expected the bundle object and reader to have the same number of bundle modules")
 	}
 
 	d.Stop(ctx)