fix test

segmentio · Aug 9, 2023 · 2545d47 · 2545d47
1 parent ffe9028
commit 2545d47
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/pkg/reflector/download.go b/pkg/reflector/download.go
@@ -53,7 +53,7 @@ func (d *S3Downloader) DownloadTo(w io.Writer) (n int64, err error) {
 	defer os.Remove(file.Name())
 
 	start := time.Now()
-	numBytes, err := downloader.Download(context.TODO(), file, &s3.GetObjectInput{
+	numBytes, err := downloader.Download(context.Background(), file, &s3.GetObjectInput{
 		Bucket: aws.String(d.Bucket),
 		Key:    aws.String(d.Key),
 	})
@@ -76,6 +76,11 @@ func (d *S3Downloader) DownloadTo(w io.Writer) (n int64, err error) {
 		if err != nil {
 			return n, errors.Wrap(err, "unzip snapshot")
 		}
+	} else {
+		n, err = io.Copy(w, file)
+		if err != nil {
+			return n, errors.Wrap(err, "copy snapshot")
+		}
 	}
 
 	events.Log("LDB inflated %d -> %d bytes", numBytes, n)

diff --git a/pkg/reflector/download_test.go b/pkg/reflector/download_test.go
@@ -38,7 +38,8 @@ func TestS3DownloadErrors(t *testing.T) {
 			s3Client: func() reflector.S3Client {
 				f := &fakes.FakeS3Client{}
 				f.GetObjectReturns(&s3.GetObjectOutput{
-					Body: ioutil.NopCloser(strings.NewReader("data")),
+					Body:          ioutil.NopCloser(strings.NewReader("data")),
+					ContentLength: 4,
 				}, nil)
 				return f
 			},