From ad37fbb762a464aecc68d5a9e5706f8552fc9d6c Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong.manhle.vn@gmail.com>
Date: Wed, 30 Oct 2019 18:46:43 +0700
Subject: [PATCH] cmd: add --summary-export flag

Add --summary-export flag for "k6 run", the parameter is output file to
write report summary in JSON format.

Empty string means output is ignored.

Close #355
---
 cmd/config.go  | 20 +++++++++++++++-----
 cmd/run.go     | 33 +++++++++++++++++++++++++++------
 core/engine.go | 13 +++++++------
 3 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/cmd/config.go b/cmd/config.go
index 9f001e37565..120f3331330 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -56,17 +56,23 @@ func configFlagSet() *pflag.FlagSet {
 	flags.Bool("no-usage-report", false, "don't send anonymous stats to the developers")
 	flags.Bool("no-thresholds", false, "don't run thresholds")
 	flags.Bool("no-summary", false, "don't show the summary at the end of the test")
+	flags.String(
+		"summary-export",
+		"",
+		"output summary to file in json format, empty string means output is ignored",
+	)
 	return flags
 }
 
 type Config struct {
 	lib.Options
 
-	Out           []string  `json:"out" envconfig:"K6_OUT"`
-	Linger        null.Bool `json:"linger" envconfig:"K6_LINGER"`
-	NoUsageReport null.Bool `json:"noUsageReport" envconfig:"K6_NO_USAGE_REPORT"`
-	NoThresholds  null.Bool `json:"noThresholds" envconfig:"K6_NO_THRESHOLDS"`
-	NoSummary     null.Bool `json:"noSummary" envconfig:"K6_NO_SUMMARY"`
+	Out           []string    `json:"out" envconfig:"K6_OUT"`
+	Linger        null.Bool   `json:"linger" envconfig:"K6_LINGER"`
+	NoUsageReport null.Bool   `json:"noUsageReport" envconfig:"K6_NO_USAGE_REPORT"`
+	NoThresholds  null.Bool   `json:"noThresholds" envconfig:"K6_NO_THRESHOLDS"`
+	NoSummary     null.Bool   `json:"noSummary" envconfig:"K6_NO_SUMMARY"`
+	SummaryExport null.String `json:"summaryExport" envconfig:"K6_SUMMARY_EXPORT"`
 
 	Collectors struct {
 		InfluxDB influxdb.Config `json:"influxdb"`
@@ -95,6 +101,9 @@ func (c Config) Apply(cfg Config) Config {
 	if cfg.NoSummary.Valid {
 		c.NoSummary = cfg.NoSummary
 	}
+	if cfg.SummaryExport.Valid {
+		c.SummaryExport = cfg.SummaryExport
+	}
 	c.Collectors.InfluxDB = c.Collectors.InfluxDB.Apply(cfg.Collectors.InfluxDB)
 	c.Collectors.Cloud = c.Collectors.Cloud.Apply(cfg.Collectors.Cloud)
 	c.Collectors.Kafka = c.Collectors.Kafka.Apply(cfg.Collectors.Kafka)
@@ -121,6 +130,7 @@ func getConfig(flags *pflag.FlagSet) (Config, error) {
 		NoUsageReport: getNullBool(flags, "no-usage-report"),
 		NoThresholds:  getNullBool(flags, "no-thresholds"),
 		NoSummary:     getNullBool(flags, "no-summary"),
+		SummaryExport: getNullString(flags, "summary-export"),
 	}, nil
 }
 
diff --git a/cmd/run.go b/cmd/run.go
index 097e4f3b0fa..db8d9f4b8e3 100644
--- a/cmd/run.go
+++ b/cmd/run.go
@@ -205,6 +205,9 @@ a commandline interface for interacting with it.`,
 		if conf.NoSummary.Valid {
 			engine.NoSummary = conf.NoSummary.Bool
 		}
+		if conf.SummaryExport.Valid {
+			engine.SummaryExport = conf.SummaryExport.String != ""
+		}
 
 		// Create a collector and assign it to the engine if requested.
 		fprintf(stdout, "%s   collector\r", initBar.String())
@@ -437,21 +440,39 @@ a commandline interface for interacting with it.`,
 			logrus.Warn("No data generated, because no script iterations finished, consider making the test duration longer")
 		}
 
+		data := ui.SummaryData{
+			Metrics:   engine.Metrics,
+			RootGroup: engine.Executor.GetRunner().GetDefaultGroup(),
+			Time:      engine.Executor.GetTime(),
+			TimeUnit:  conf.Options.SummaryTimeUnit.String,
+		}
 		// Print the end-of-test summary.
 		if !conf.NoSummary.Bool {
 			fprintf(stdout, "\n")
 
 			s := ui.NewSummary(conf.SummaryTrendStats)
-			s.SummarizeMetrics(stdout, "", ui.SummaryData{
-				Metrics:   engine.Metrics,
-				RootGroup: engine.Executor.GetRunner().GetDefaultGroup(),
-				Time:      engine.Executor.GetTime(),
-				TimeUnit:  conf.Options.SummaryTimeUnit.String,
-			})
+			s.SummarizeMetrics(stdout, "", data)
 
 			fprintf(stdout, "\n")
 		}
 
+		if conf.SummaryExport.Valid && conf.SummaryExport.String != "" {
+			f, err := os.Create(conf.SummaryExport.String)
+			if err != nil {
+				logrus.Error("failed to create summary export file")
+			} else {
+				defer func() {
+					if err := f.Close(); err != nil {
+						logrus.WithError(err).Fatal("failed to close summary output file")
+					}
+				}()
+				s := ui.NewSummary(conf.SummaryTrendStats)
+				if err := s.SummarizeMetricsJSON(f, data); err != nil {
+					logrus.WithError(err).Fatal("failed to make summary export")
+				}
+			}
+		}
+
 		if conf.Linger.Bool {
 			logrus.Info("Linger set; waiting for Ctrl+C...")
 			<-sigC
diff --git a/core/engine.go b/core/engine.go
index 5189a30c23b..fd4540b8bbc 100644
--- a/core/engine.go
+++ b/core/engine.go
@@ -50,11 +50,12 @@ const (
 type Engine struct {
 	runLock sync.Mutex
 
-	Executor     lib.Executor
-	Options      lib.Options
-	Collectors   []lib.Collector
-	NoThresholds bool
-	NoSummary    bool
+	Executor      lib.Executor
+	Options       lib.Options
+	Collectors    []lib.Collector
+	NoThresholds  bool
+	NoSummary     bool
+	SummaryExport bool
 
 	logger *logrus.Logger
 
@@ -386,7 +387,7 @@ func (e *Engine) processSamples(sampleContainers []stats.SampleContainer) {
 	defer e.MetricsLock.Unlock()
 
 	// TODO: run this and the below code in goroutines?
-	if !(e.NoSummary && e.NoThresholds) {
+	if !(e.NoSummary && e.NoThresholds && !e.SummaryExport) {
 		e.processSamplesForMetrics(sampleContainers)
 	}