Start adding compressed download testing

download/downloader_test.go

@@ -2,6 +2,7 @@ package download
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"errors"
 	"fmt"
@@ -9,44 +10,27 @@ import (
 	"testing"
 )
 
-type mockStorageClient struct {
-	data  []byte
-	error error
-}
-
-func (m *mockStorageClient) Download(ctx context.Context, writer io.WriterAt) error {
-	if m.error != nil {
-		return m.error
-	}
-	n, err := writer.WriteAt(m.data, 0)
-	if n != len(m.data) || err != nil {
-		return fmt.Errorf("failed to write data: %w", err)
-	}
-	return nil
-}
-
-func (m *mockStorageClient) String() string {
-	return "mockStorageClient"
-}
-
 func TestDownloader_Do(t *testing.T) {
 	tests := []struct {
 		name           string
 		mockClientData []byte
-		mockClientError    error
-		expectedOutputData []byte
-		expectError        bool
+		compress       bool
+		expectError    error
 	}{
 		{
-			name:               "Successful Download",
+			name:           "Successful download",
+			mockClientData: []byte("test data"),
+			expectError:    nil,
+		},
+		{
+			name:           "Successful download of compressed data",
 			mockClientData: []byte("test data"),
-			expectedOutputData: []byte("test data"),
-			expectError:        false,
+			compress:       false,
+			expectError:    nil,
 		},
 		{
-			name:            "Download Error",
-			mockClientError: errors.New("download error"),
-			expectError:     true,
+			name:        "Download error",
+			expectError: errors.New("download error"),
 		},
 	}
 
@@ -54,55 +38,92 @@ func TestDownloader_Do(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			mockClient := &mockStorageClient{
 				data:  tt.mockClientData,
-				error: tt.mockClientError,
+				error: tt.expectError,
 			}
-			downloader := download.NewDownloader(mockClient)
+			if tt.compress {
+				mockClient.Compress()
+			}
+			downloader := NewDownloader(mockClient)
 
-			outputBuffer := bytes.NewBuffer(make([]byte, 0, len(tt.expectedOutputData)))
-			err := downloader.Do(context.Background(), outputBuffer)
+			f := newmockWriterAt(len(tt.mockClientData))
 
-			if tt.expectError {
+			err := downloader.Do(context.Background(), f)
+			if tt.expectError != nil {
 				if err == nil {
 					t.Errorf("Expected error, but got none")
 				}
+				if err.Error() != tt.expectError.Error() {
+					t.Errorf("Expected error %v, but got %v", tt.expectError, err)
+				}
 			} else {
 				if err != nil {
 					t.Errorf("Unexpected error: %v", err)
 				}
-				if !bytes.Equal(tt.expectedOutputData, outputBuffer.Bytes()) {
-					t.Errorf("Expected output data %v, but got %v", tt.expectedOutputData, outputBuffer.Bytes())
+				if !bytes.Equal(tt.mockClientData, f.Bytes()) {
+					t.Errorf("Expected output data %v, but got %v", tt.mockClientData, f.Bytes())
 				}
 			}
 		})
 	}
 }
 
-func TestDownloader_Stats(t *testing.T) {
-	mockClient := &mockStorageClient{
-		data: []byte("test data"),
+type mockStorageClient struct {
+	data  []byte
+	error error
+}
+
+func (m *mockStorageClient) Download(ctx context.Context, writer io.WriterAt) error {
+	if m.error != nil {
+		return m.error
+	}
+	n, err := writer.WriteAt(m.data, 0)
+	if n != len(m.data) || err != nil {
+		return fmt.Errorf("failed to write data: %w", err)
+	}
+	return nil
 }
-	downloader := download.NewDownloader(mockClient)
-	buffer := bytes.NewBuffer(make([]byte, 0, len(mockClient.data)))
 
-	download.ResetStats()
+func (m *mockStorageClient) Compress() error {
+	var compressedData bytes.Buffer
+	gzipWriter := gzip.NewWriter(&compressedData)
 
-	err := downloader.Do(context.Background(), buffer)
+	_, err := gzipWriter.Write(m.data)
 	if err != nil {
-		t.Errorf("Unexpected error: %v", err)
+		return err
+	}
+
+	err = gzipWriter.Close()
+	if err != nil {
+		return err
+	}
+
+	m.data = compressedData.Bytes()
+	return nil
+}
+
+func (m *mockStorageClient) String() string {
+	return "mockStorageClient"
 }
 
-	numDownloadsOK := download.GetStats().Get(download.NumDownloadsOK)
-	if numDownloadsOK != 1 {
-		t.Errorf("Expected NumDownloadsOK to be 1, but got %d", numDownloadsOK)
+type mockWriterAt struct {
+	data []byte
 }
 
-	numDownloadsFail := download.GetStats().Get(download.NumDownloadsFail)
-	if numDownloadsFail != 0 {
-		t.Errorf("Expected NumDownloadsFail to be 0, but got %d", numDownloadsFail)
+func newmockWriterAt(size int) *mockWriterAt {
+	return &mockWriterAt{
+		data: make([]byte, size),
+	}
 }
 
-	numDownloadBytes := download.GetStats().Get(download.NumDownloadBytes)
-	if numDownloadBytes != int64(len(mockClient.data)) {
-		t.Errorf("Expected NumDownloadBytes to be %d, but got %d", len(mockClient.data), numDownloadBytes)
+func (s *mockWriterAt) WriteAt(p []byte, off int64) (n int, err error) {
+	if off < 0 || int(off) > len(s.data) {
+		return 0, errors.New("invalid offset")
 	}
+
+	n = copy(s.data[off:], p)
+	return n, nil
+}
+
+func (s *mockWriterAt) Bytes() []byte {
+	return s.data
 }