From 8ce65c138e22d80307c47a3102b01ce547520377 Mon Sep 17 00:00:00 2001
From: Ivan <2103732+codebien@users.noreply.github.com>
Date: Sun, 7 Mar 2021 12:24:11 +0100
Subject: [PATCH] Thresholds flush for the json output

The JSON Output implements WithThresholds
so the Engine can set the thresholds invoking SetThresholds.

The flush logic sets the Metric's Thresholds
to write them in the json output.

Closes #1052
---
 output/json/json.go      | 13 ++++++++++++-
 output/json/json_test.go | 21 +++++++++++++++++++--
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/output/json/json.go b/output/json/json.go
index f2d8ea36806..49591d2f951 100644
--- a/output/json/json.go
+++ b/output/json/json.go
@@ -32,7 +32,7 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-// TODO: add option for emitting proper JSON files (https://github.com/loadimpact/k6/issues/1052)
+// TODO: add option for emitting proper JSON files (https://github.com/loadimpact/k6/issues/737)
 const flushPeriod = 200 * time.Millisecond // TODO: make this configurable
 
 // Output funnels all passed metrics to an (optionally gzipped) JSON file.
@@ -47,6 +47,7 @@ type Output struct {
 	encoder     *stdlibjson.Encoder
 	closeFn     func() error
 	seenMetrics map[string]struct{}
+	thresholds  map[string][]*stats.Threshold
 }
 
 // New returns a new JSON output.
@@ -118,6 +119,15 @@ func (o *Output) Stop() error {
 	return o.closeFn()
 }
 
+// SetThresholds receives the thresholds before the output is Start()-ed.
+func (o *Output) SetThresholds(thresholds map[string]stats.Thresholds) {
+	ths := make(map[string][]*stats.Threshold)
+	for name, t := range thresholds {
+		ths[name] = append(ths[name], t.Thresholds...)
+	}
+	o.thresholds = ths
+}
+
 func (o *Output) flushMetrics() {
 	samples := o.GetBufferedSamples()
 	start := time.Now()
@@ -127,6 +137,7 @@ func (o *Output) flushMetrics() {
 		count += len(samples)
 		for _, sample := range samples {
 			sample := sample
+			sample.Metric.Thresholds.Thresholds = o.thresholds[sample.Metric.Name]
 			o.handleMetric(sample.Metric)
 			err := o.encoder.Encode(WrapSample(sample))
 			if err != nil {
diff --git a/output/json/json_test.go b/output/json/json_test.go
index da104f95261..e273a3871aa 100644
--- a/output/json/json_test.go
+++ b/output/json/json_test.go
@@ -71,9 +71,8 @@ func generateTestMetricSamples(t *testing.T) ([]stats.SampleContainer, func(io.R
 		}, Time: time2, Tags: connTags},
 		stats.Sample{Time: time3, Metric: metric2, Value: float64(5), Tags: stats.NewSampleTags(map[string]string{"tag3": "val3"})},
 	}
-	// TODO: fix JSON thresholds (https://github.com/loadimpact/k6/issues/1052)
 	expected := []string{
-		`{"type":"Metric","data":{"name":"my_metric1","type":"gauge","contains":"default","tainted":null,"thresholds":[],"submetrics":null,"sub":{"name":"","parent":"","suffix":"","tags":null}},"metric":"my_metric1"}`,
+		`{"type":"Metric","data":{"name":"my_metric1","type":"gauge","contains":"default","tainted":null,"thresholds":["rate<0.01", "p(99)<250"],"submetrics":null,"sub":{"name":"","parent":"","suffix":"","tags":null}},"metric":"my_metric1"}`,
 		`{"type":"Point","data":{"time":"2021-02-24T13:37:10Z","value":1,"tags":{"tag1":"val1"}},"metric":"my_metric1"}`,
 		`{"type":"Point","data":{"time":"2021-02-24T13:37:10Z","value":2,"tags":{"tag2":"val2"}},"metric":"my_metric1"}`,
 		`{"type":"Metric","data":{"name":"my_metric2","type":"counter","contains":"data","tainted":null,"thresholds":[],"submetrics":null,"sub":{"name":"","parent":"","suffix":"","tags":null}},"metric":"my_metric2"}`,
@@ -94,6 +93,8 @@ func TestJsonOutputStdout(t *testing.T) {
 		StdOut: stdout,
 	})
 	require.NoError(t, err)
+
+	setThresholds(t, out)
 	require.NoError(t, out.Start())
 
 	samples, validateResults := generateTestMetricSamples(t)
@@ -130,6 +131,8 @@ func TestJsonOutputFile(t *testing.T) {
 		ConfigArgument: "/json-output",
 	})
 	require.NoError(t, err)
+
+	setThresholds(t, out)
 	require.NoError(t, out.Start())
 
 	samples, validateResults := generateTestMetricSamples(t)
@@ -156,6 +159,8 @@ func TestJsonOutputFileGzipped(t *testing.T) {
 		ConfigArgument: "/json-output.gz",
 	})
 	require.NoError(t, err)
+
+	setThresholds(t, out)
 	require.NoError(t, out.Start())
 
 	samples, validateResults := generateTestMetricSamples(t)
@@ -185,3 +190,15 @@ func TestWrapMetricWithMetricPointer(t *testing.T) {
 	out := wrapMetric(&stats.Metric{})
 	assert.NotEqual(t, out, (*Envelope)(nil))
 }
+
+func setThresholds(t *testing.T, out output.Output) {
+	t.Helper()
+
+	jout, ok := out.(*Output)
+	require.True(t, ok)
+
+	ts, err := stats.NewThresholds([]string{"rate<0.01", "p(99)<250"})
+	require.NoError(t, err)
+
+	jout.SetThresholds(map[string]stats.Thresholds{"my_metric1": ts})
+}