From dc9ca9c23b98e180c8f1ba836baa7876e4b58f9a Mon Sep 17 00:00:00 2001
From: Zhenghao Lou <rhzx3519@gmail.com>
Date: Wed, 29 Mar 2023 21:23:37 +1100
Subject: [PATCH] feat: helm test --logs, support outputing logs only from
 failed pods, succeeded pods, or all pods.

helm test --logs all (default) // output logs of all pods
helm test --logs failed  // only output logs of failed pods
helm test --logs succeeded // only output logs of succeeded pods

Signed-off-by: Zhenghao Lou <rhzx3519@gmail.com>
---
 cmd/helm/release_testing.go   | 14 ++++++++++----
 pkg/action/release_testing.go | 22 +++++++++++++++++++---
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/cmd/helm/release_testing.go b/cmd/helm/release_testing.go
index d9b8fa8c9..71ec09ae6 100644
--- a/cmd/helm/release_testing.go
+++ b/cmd/helm/release_testing.go
@@ -40,7 +40,7 @@ The tests to be run are defined in the chart that was installed.
 func newReleaseTestCmd(cfg *action.Configuration, out io.Writer) *cobra.Command {
 	client := action.NewReleaseTesting(cfg)
 	var outfmt = output.Table
-	var outputLogs bool
+	var outputLogs string
 	var filter []string
 
 	cmd := &cobra.Command{
@@ -76,10 +76,16 @@ func newReleaseTestCmd(cfg *action.Configuration, out io.Writer) *cobra.Command
 				return err
 			}
 
-			if outputLogs {
+			if outputLogs != "" {
+				podLogFilter := action.PodAll
+				if outputLogs == string(action.PodSucceeded) {
+					podLogFilter = action.PodSucceeded
+				} else if outputLogs == string(action.PodFailed) {
+					podLogFilter = action.PodFailed
+				}
 				// Print a newline to stdout to separate the output
 				fmt.Fprintln(out)
-				if err := client.GetPodLogs(out, rel); err != nil {
+				if err := client.GetPodLogs(out, rel, podLogFilter); err != nil {
 					return err
 				}
 			}
@@ -90,7 +96,7 @@ func newReleaseTestCmd(cfg *action.Configuration, out io.Writer) *cobra.Command
 
 	f := cmd.Flags()
 	f.DurationVar(&client.Timeout, "timeout", 300*time.Second, "time to wait for any individual Kubernetes operation (like Jobs for hooks)")
-	f.BoolVar(&outputLogs, "logs", false, "dump the logs from test pods (this runs after all tests are complete, but before any cleanup)")
+	f.StringVar(&outputLogs, "logs", "", "dump the logs from test pods (this runs after all tests are complete, but before any cleanup), options include [\"all\", \"succeeded\", \"failed\"]")
 	f.StringSliceVar(&filter, "filter", []string{}, "specify tests by attribute (currently \"name\") using attribute=value syntax or '!attribute=value' to exclude a test (can specify multiple or separate values with commas: name=test1,name=test2)")
 
 	return cmd
diff --git a/pkg/action/release_testing.go b/pkg/action/release_testing.go
index ecaeaf59f..81525ae0c 100644
--- a/pkg/action/release_testing.go
+++ b/pkg/action/release_testing.go
@@ -23,10 +23,17 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
-	v1 "k8s.io/api/core/v1"
-
 	"helm.sh/helm/v3/pkg/chartutil"
 	"helm.sh/helm/v3/pkg/release"
+	v1 "k8s.io/api/core/v1"
+)
+
+type PodLogFilter string
+
+const (
+	PodAll       PodLogFilter = "all"
+	PodSucceeded PodLogFilter = "succeeded"
+	PodFailed    PodLogFilter = "failed"
 )
 
 // ReleaseTesting is the action for testing a release.
@@ -101,7 +108,7 @@ func (r *ReleaseTesting) Run(name string) (*release.Release, error) {
 // GetPodLogs will write the logs for all test pods in the given release into
 // the given writer. These can be immediately output to the user or captured for
 // other uses
-func (r *ReleaseTesting) GetPodLogs(out io.Writer, rel *release.Release) error {
+func (r *ReleaseTesting) GetPodLogs(out io.Writer, rel *release.Release, filter PodLogFilter) error {
 	client, err := r.cfg.KubernetesClientSet()
 	if err != nil {
 		return errors.Wrap(err, "unable to get kubernetes client to fetch pod logs")
@@ -110,6 +117,15 @@ func (r *ReleaseTesting) GetPodLogs(out io.Writer, rel *release.Release) error {
 	for _, h := range rel.Hooks {
 		for _, e := range h.Events {
 			if e == release.HookTest {
+				// only print logs of failed pods
+				if filter == PodFailed && h.LastRun.Phase != release.HookPhaseFailed {
+					continue
+				}
+				// only print logs of succeeded pods
+				if filter == PodSucceeded && h.LastRun.Phase != release.HookPhaseSucceeded {
+					continue
+				}
+
 				req := client.CoreV1().Pods(r.Namespace).GetLogs(h.Name, &v1.PodLogOptions{})
 				logReader, err := req.Stream(context.Background())
 				if err != nil {