Tests added for download function

Author: shiv-am0

plugin.go

@@ -141,74 +141,7 @@ func (p *Plugin) Exec() error {
 	}
 
 	if p.Download {
-		targetDir := strings.TrimPrefix(filepath.ToSlash(p.Target), "/")
-		log.WithFields(log.Fields{
-			"bucket": p.Bucket,
-			"dir":    targetDir,
-		}).Info("Listing S3 directory")
-
-		list, err := client.ListObjectsV2(&s3.ListObjectsV2Input{
-			Bucket: &p.Bucket,
-			Prefix: &targetDir,
-		})
-		if err != nil {
-			log.WithFields(log.Fields{
-				"error":  err,
-				"bucket": p.Bucket,
-				"dir":    targetDir,
-			}).Error("Cannot list S3 directory")
-			return err
-		}
-
-		g := errgroup.Group{}
-
-		for _, item := range list.Contents {
-			log.WithFields(log.Fields{
-				"bucket": p.Bucket,
-				"key":    *item.Key,
-			}).Info("Getting S3 object")
-
-			item := item
-			g.Go(func() error {
-				obj, err := client.GetObject(&s3.GetObjectInput{
-					Bucket: &p.Bucket,
-					Key:    item.Key,
-				})
-				if err != nil {
-					log.WithFields(log.Fields{
-						"error":  err,
-						"bucket": p.Bucket,
-						"key":    *item.Key,
-					}).Error("Cannot get S3 object")
-					return err
-				}
-
-				source := resolveSource(targetDir, *item.Key, p.StripPrefix)
-
-				f, err := os.Create(source)
-				if err != nil {
-					log.WithFields(log.Fields{
-						"error": err,
-						"file":  source,
-					}).Error("Problem opening file for writing")
-					return err
-				}
-				defer f.Close()
-
-				_, err = io.Copy(f, obj.Body)
-				if err != nil {
-					log.WithFields(log.Fields{
-						"error": err,
-						"file":  source,
-					}).Error("Failed to write file")
-					return err
-				}
-
-				return nil
-			})
-		}
-
-		return g.Wait()
+		downloadFromS3(p, client)
 	}
 
 	// find the bucket
@@ -423,3 +356,79 @@ func isDir(source string, matches []string) bool {
 	}
 	return false
 }
+
+func resolveTargetDir(s string) string {
+	res := strings.TrimPrefix(filepath.ToSlash(s), "/")
+	return res
+}
+
+func downloadFromS3(p *Plugin, client *s3.S3) error {
+	targetDir := resolveTargetDir(p.Target)
+	log.WithFields(log.Fields{
+		"bucket": p.Bucket,
+		"dir":    targetDir,
+	}).Info("Listing S3 directory")
+
+	list, err := client.ListObjectsV2(&s3.ListObjectsV2Input{
+		Bucket: &p.Bucket,
+		Prefix: &targetDir,
+	})
+	if err != nil {
+		log.WithFields(log.Fields{
+			"error":  err,
+			"bucket": p.Bucket,
+			"dir":    targetDir,
+		}).Error("Cannot list S3 directory")
+		return err
+	}
+
+	g := errgroup.Group{}
+
+	for _, item := range list.Contents {
+		log.WithFields(log.Fields{
+			"bucket": p.Bucket,
+			"key":    *item.Key,
+		}).Info("Getting S3 object")
+
+		item := item
+		g.Go(func() error {
+			obj, err := client.GetObject(&s3.GetObjectInput{
+				Bucket: &p.Bucket,
+				Key:    item.Key,
+			})
+			if err != nil {
+				log.WithFields(log.Fields{
+					"error":  err,
+					"bucket": p.Bucket,
+					"key":    *item.Key,
+				}).Error("Cannot get S3 object")
+				return err
+			}
+
+			source := resolveSource(targetDir, *item.Key, p.StripPrefix)
+
+			f, err := os.Create(source)
+			if err != nil {
+				log.WithFields(log.Fields{
+					"error": err,
+					"file":  source,
+				}).Error("Problem opening file for writing")
+				return err
+			}
+			defer f.Close()
+
+			_, err = io.Copy(f, obj.Body)
+			if err != nil {
+				log.WithFields(log.Fields{
+					"error": err,
+					"file":  source,
+				}).Error("Failed to write file")
+				return err
+			}
+
+			return nil
+		})
+	}
+
+	return g.Wait()
+}

plugin_unix_test.go

@@ -44,3 +44,38 @@ func TestResolveUnixKey(t *testing.T) {
 		}
 	}
 }
+
+func TestResolveTargetDir(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected string
+	}{
+		{
+			input:    "example-string",
+			expected: "example-string",
+		},
+		{
+			input:    "/path/to/file",
+			expected: "path/to/file",
+		},
+		{
+			input:    "12345",
+			expected: "12345",
+		},
+		{
+			input:    "/root/directory",
+			expected: "root/directory",
+		},
+		{
+			input:    "no_slash",
+			expected: "no_slash",
+		},
+	}
+
+	for _, tc := range tests {
+		result := resolveTargetDir(tc.input)
+		if result != tc.expected {
+			t.Errorf("Expected: %s, Got: %s", tc.expected, result)
+		}
+	}
+}