feat: Log stats of artifact pull command

Author: hamir-suspect

cmd/pull.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"fmt"
 	errutil "github.com/semaphoreci/artifact/pkg/errors"
 	"github.com/semaphoreci/artifact/pkg/files"
 	"github.com/semaphoreci/artifact/pkg/hub"
@@ -18,7 +19,7 @@ artifact push. With artifact pull you can download them to the current directory
 to use them in a later phase, debug, or getting the results.`,
 }
 
-func runPullForCategory(cmd *cobra.Command, args []string, resolver *files.PathResolver) (*files.ResolvedPath, error) {
+func runPullForCategory(cmd *cobra.Command, args []string, resolver *files.PathResolver) (*files.ResolvedPath, *storage.PullStats, error) {
 	destinationOverride, err := cmd.Flags().GetString("destination")
 	errutil.Check(err)
 
@@ -49,7 +50,7 @@ func NewPullJobCmd() *cobra.Command {
 			resolver, err := files.NewPathResolver(files.ResourceTypeJob, jobId)
 			errutil.Check(err)
 
-			paths, err := runPullForCategory(cmd, args, resolver)
+			paths, stats, err := runPullForCategory(cmd, args, resolver)
 			if err != nil {
 				log.Errorf("Error pulling artifact: %v\n", err)
 				log.Error("Please check if the artifact you are trying to pull exists.\n")
@@ -60,6 +61,7 @@ func NewPullJobCmd() *cobra.Command {
 			log.Info("Successfully pulled artifact for current job.\n")
 			log.Infof("* Remote source: '%s'.\n", paths.Source)
 			log.Infof("* Local destination: '%s'.\n", paths.Destination)
+			log.Infof("Pulled %d files. Total of %s\n", stats.FileCount, formatBytes(stats.TotalSize))
 		},
 	}
 
@@ -83,7 +85,7 @@ func NewPullWorkflowCmd() *cobra.Command {
 			resolver, err := files.NewPathResolver(files.ResourceTypeWorkflow, workflowId)
 			errutil.Check(err)
 
-			paths, err := runPullForCategory(cmd, args, resolver)
+			paths, stats, err := runPullForCategory(cmd, args, resolver)
 			if err != nil {
 				log.Errorf("Error pulling artifact: %v\n", err)
 				log.Error("Please check if the artifact you are trying to pull exists.\n")
@@ -94,6 +96,7 @@ func NewPullWorkflowCmd() *cobra.Command {
 			log.Info("Successfully pulled artifact for current workflow.\n")
 			log.Infof("* Remote source: '%s'.\n", paths.Source)
 			log.Infof("* Local destination: '%s'.\n", paths.Destination)
+			log.Infof("Pulled %d files. Total of %s\n", stats.FileCount, formatBytes(stats.TotalSize))
 		},
 	}
 
@@ -117,7 +120,7 @@ func NewPullProjectCmd() *cobra.Command {
 			resolver, err := files.NewPathResolver(files.ResourceTypeProject, projectId)
 			errutil.Check(err)
 
-			paths, err := runPullForCategory(cmd, args, resolver)
+			paths, stats, err := runPullForCategory(cmd, args, resolver)
 			if err != nil {
 				log.Errorf("Error pulling artifact: %v\n", err)
 				log.Error("Please check if the artifact you are trying to pull exists.\n")
@@ -128,6 +131,7 @@ func NewPullProjectCmd() *cobra.Command {
 			log.Info("Successfully pulled artifact for current project.\n")
 			log.Infof("* Remote source: '%s'.\n", paths.Source)
 			log.Infof("* Local destination: '%s'.\n", paths.Destination)
+			log.Infof("Pulled %d files. Total of %s\n", stats.FileCount, formatBytes(stats.TotalSize))
 		},
 	}
 
@@ -137,6 +141,20 @@ func NewPullProjectCmd() *cobra.Command {
 	return cmd
 }
 
+// formatBytes converts bytes to human readable format
+func formatBytes(bytes int64) string {
+	const unit = 1024
+	if bytes < unit {
+		return fmt.Sprintf("%d B", bytes)
+	}
+	div, exp := int64(unit), 0
+	for n := bytes / unit; n >= unit; n /= unit {
+		div *= unit
+		exp++
+	}
+	return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
+}
+
 func init() {
 	rootCmd.AddCommand(pullCmd)
 	pullCmd.AddCommand(NewPullJobCmd())

cmd/pull_test.go

@@ -172,3 +172,27 @@ func assertFileDoesNotExist(t *testing.T, fileName string) {
 	_, err := os.Stat(fileName)
 	assert.True(t, os.IsNotExist(err))
 }
+
+func Test__formatBytes(t *testing.T) {
+	testCases := []struct {
+		name     string
+		bytes    int64
+		expected string
+	}{
+		{"zero bytes", 0, "0 B"},
+		{"small bytes", 512, "512 B"},
+		{"exactly 1KB", 1024, "1.0 KB"},
+		{"1.5KB", 1536, "1.5 KB"},
+		{"exactly 1MB", 1024*1024, "1.0 MB"},
+		{"498MB", 498*1024*1024, "498.0 MB"},
+		{"1.2GB", int64(1288490188), "1.2 GB"},
+		{"large size", 5*1024*1024*1024*1024, "5.0 TB"},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := formatBytes(tc.bytes)
+			assert.Equal(t, tc.expected, result)
+		})
+	}
+}

docker-compose.yml

@@ -7,6 +7,8 @@ services:
     tty: true
     command: "sleep 0"
     container_name: 'artifact-cli'
+    environment:
+      - HTTP_PROXY=localhost:3000
     volumes:
       - go-pkg-cache:/go
       - .:/app

pkg/hub/hub_test.go

@@ -99,6 +99,15 @@ func Test__GenerateSignedURL(t *testing.T) {
 			assert.Contains(t, err.Error(), "request did not return a non-5xx response")
 		}
 	})
+
+	t.Run("Check if http_proxy enviroment variable is used", func(t *testing.T) {
+		noOfCalls := 0
+		mockArtifactHubServer := generateMockServer(&noOfCalls, 500, []byte("{}"))
+		defer mockArtifactHubServer.Close()
+
+		generateSignedURLsHelper(mockArtifactHubServer.URL)
+
+	})
 }
 
 func generateSignedURLsHelper(url string) (*GenerateSignedURLsResponse, error) {

pkg/storage/pull.go

@@ -17,10 +17,15 @@ type PullOptions struct {
 	Force               bool
 }
 
-func Pull(hubClient *hub.Client, resolver *files.PathResolver, options PullOptions) (*files.ResolvedPath, error) {
+type PullStats struct {
+	FileCount int
+	TotalSize int64
+}
+
+func Pull(hubClient *hub.Client, resolver *files.PathResolver, options PullOptions) (*files.ResolvedPath, *PullStats, error) {
 	paths, err := resolver.Resolve(files.OperationPull, options.SourcePath, options.DestinationOverride)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	log.Debug("Pulling...\n")
@@ -30,15 +35,20 @@ func Pull(hubClient *hub.Client, resolver *files.PathResolver, options PullOptio
 
 	response, err := hubClient.GenerateSignedURLs([]string{paths.Source}, hub.GenerateSignedURLsRequestPULL)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	artifacts, err := buildArtifacts(response.Urls, paths, options.Force)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
+	}
+
+	stats, err := doPull(options.Force, artifacts, response.Urls)
+	if err != nil {
+		return nil, nil, err
 	}
 
-	return paths, doPull(options.Force, artifacts, response.Urls)
+	return paths, stats, nil
 }
 
 func buildArtifacts(signedURLs []*api.SignedURL, paths *files.ResolvedPath, force bool) ([]*api.Artifact, error) {
@@ -68,16 +78,23 @@ func buildArtifacts(signedURLs []*api.SignedURL, paths *files.ResolvedPath, forc
 	return artifacts, nil
 }
 
-func doPull(force bool, artifacts []*api.Artifact, signedURLs []*api.SignedURL) error {
+func doPull(force bool, artifacts []*api.Artifact, signedURLs []*api.SignedURL) (*PullStats, error) {
 	client := newHTTPClient()
+	stats := &PullStats{}
 
 	for _, artifact := range artifacts {
 		for _, signedURL := range artifact.URLs {
 			if err := signedURL.Follow(client, artifact); err != nil {
-				return err
+				return nil, err
+			}
+			
+			// Get file size after successful download
+			if fileInfo, err := os.Stat(artifact.LocalPath); err == nil {
+				stats.FileCount++
+				stats.TotalSize += fileInfo.Size()
 			}
 		}
 	}
 
-	return nil
+	return stats, nil
 }

pkg/storage/pull_test.go

@@ -0,0 +1,102 @@
+package storage
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/semaphoreci/artifact/pkg/api"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test__doPull_Stats(t *testing.T) {
+	// Create temporary directory for test files
+	tempDir, err := ioutil.TempDir("", "pull_test")
+	require.NoError(t, err)
+	defer os.RemoveAll(tempDir)
+
+	// Create test artifacts with known sizes
+	testFiles := []struct {
+		name    string
+		content string
+		size    int64
+	}{
+		{"file1.txt", "hello world", 11},
+		{"file2.txt", "test content here", 17},
+		{"file3.txt", "a", 1},
+	}
+
+	artifacts := []*api.Artifact{}
+	for _, tf := range testFiles {
+		localPath := filepath.Join(tempDir, tf.name)
+		artifacts = append(artifacts, &api.Artifact{
+			RemotePath: tf.name,
+			LocalPath:  localPath,
+			URLs:       []*api.SignedURL{}, // Empty for this test
+		})
+
+		// Pre-create the files to simulate successful downloads
+		err := ioutil.WriteFile(localPath, []byte(tf.content), 0644)
+		require.NoError(t, err)
+	}
+
+	// Mock the doPull function to skip actual HTTP calls
+	// We'll test the stats collection logic by creating a modified version
+	stats := &PullStats{}
+	
+	// Simulate the stats collection that happens in doPull
+	for _, artifact := range artifacts {
+		if fileInfo, err := os.Stat(artifact.LocalPath); err == nil {
+			stats.FileCount++
+			stats.TotalSize += fileInfo.Size()
+		}
+	}
+
+	// Verify stats
+	assert.Equal(t, 3, stats.FileCount)
+	assert.Equal(t, int64(29), stats.TotalSize) // 11 + 17 + 1
+}
+
+func Test__PullStats_EmptyDirectory(t *testing.T) {
+	// Test with no files
+	stats := &PullStats{}
+	
+	assert.Equal(t, 0, stats.FileCount)
+	assert.Equal(t, int64(0), stats.TotalSize)
+}
+
+func Test__PullStats_LargeFiles(t *testing.T) {
+	// Create temporary directory for test files
+	tempDir, err := ioutil.TempDir("", "pull_large_test")
+	require.NoError(t, err)
+	defer os.RemoveAll(tempDir)
+
+	// Create a larger test file
+	largeContent := make([]byte, 1024*1024) // 1MB
+	for i := range largeContent {
+		largeContent[i] = byte(i % 256)
+	}
+
+	localPath := filepath.Join(tempDir, "large_file.bin")
+	err = ioutil.WriteFile(localPath, largeContent, 0644)
+	require.NoError(t, err)
+
+	artifact := &api.Artifact{
+		RemotePath: "large_file.bin",
+		LocalPath:  localPath,
+		URLs:       []*api.SignedURL{},
+	}
+
+	stats := &PullStats{}
+	
+	// Simulate stats collection
+	if fileInfo, err := os.Stat(artifact.LocalPath); err == nil {
+		stats.FileCount++
+		stats.TotalSize += fileInfo.Size()
+	}
+
+	assert.Equal(t, 1, stats.FileCount)
+	assert.Equal(t, int64(1024*1024), stats.TotalSize)
+}