diff --git a/download/downloader_test.go b/download/downloader_test.go index 07ee62ef..1d9c5592 100644 --- a/download/downloader_test.go +++ b/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 + name string + mockClientData []byte + compress bool + expectError error }{ { - name: "Successful Download", - mockClientData: []byte("test data"), - expectedOutputData: []byte("test data"), - expectError: false, + name: "Successful download", + mockClientData: []byte("test data"), + expectError: nil, }, { - name: "Download Error", - mockClientError: errors.New("download error"), - expectError: true, + name: "Successful download of compressed data", + mockClientData: []byte("test data"), + compress: false, + expectError: nil, + }, + { + 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, + } + if tt.compress { + mockClient.Compress() } - downloader := download.NewDownloader(mockClient) + 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 } - downloader := download.NewDownloader(mockClient) - buffer := bytes.NewBuffer(make([]byte, 0, len(mockClient.data))) + 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 +} - 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 } - numDownloadsOK := download.GetStats().Get(download.NumDownloadsOK) - if numDownloadsOK != 1 { - t.Errorf("Expected NumDownloadsOK to be 1, but got %d", numDownloadsOK) + err = gzipWriter.Close() + if err != nil { + return err } - numDownloadsFail := download.GetStats().Get(download.NumDownloadsFail) - if numDownloadsFail != 0 { - t.Errorf("Expected NumDownloadsFail to be 0, but got %d", numDownloadsFail) + m.data = compressedData.Bytes() + return nil +} + +func (m *mockStorageClient) String() string { + return "mockStorageClient" +} + +type mockWriterAt struct { + data []byte +} + +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 }