From bc779d8c79c6ef1b5a249cafbe2c5ca4e8bccb2f Mon Sep 17 00:00:00 2001
From: carla <kirkcohenc@gmail.com>
Date: Thu, 16 Jan 2020 12:07:22 +0200
Subject: [PATCH] dataset: add get threshold function to dataset

---
 dataset/dataset.go      | 26 ++++++++++++++++
 dataset/dataset_test.go | 67 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/dataset/dataset.go b/dataset/dataset.go
index d14bb39..e0b2453 100644
--- a/dataset/dataset.go
+++ b/dataset/dataset.go
@@ -180,3 +180,29 @@ func (d Dataset) GetOutliers(outlierMultiplier float64) (map[string]*OutlierResu
 
 	return outliers, nil
 }
+
+// GetThreshold returns the set of values in a dataset <= or > a given
+// threshold. The below bool is used to toggle whether we identify values
+// above or below the threshold.
+func (d Dataset) GetThreshold(thresholdValue float64, below bool) map[string]bool {
+	threshold := make(map[string]bool, len(d.rawValues()))
+
+	for label, value := range d {
+		// If we are looking for values below the threshold, check the
+		// current value then move on to the next one.
+		if below {
+			// If the value is below or equal to the threshold, we
+			// set the label's value in the map to true. Otherwise
+			// we set it to false.
+			threshold[label] = value <= thresholdValue
+			continue
+		}
+
+		// We are looking for values above the threshold. Set the
+		// label's value to true if the value is greater than the
+		// threshold.
+		threshold[label] = value > thresholdValue
+	}
+
+	return threshold
+}
diff --git a/dataset/dataset_test.go b/dataset/dataset_test.go
index 9d3c20d..6a9125b 100644
--- a/dataset/dataset_test.go
+++ b/dataset/dataset_test.go
@@ -2,6 +2,7 @@ package dataset
 
 import (
 	"fmt"
+	"reflect"
 	"testing"
 )
 
@@ -242,3 +243,69 @@ func TestIsOutlier(t *testing.T) {
 		})
 	}
 }
+
+// TestGetThreshold tests getting thresholds for a dataset.
+func TestGetThreshold(t *testing.T) {
+	tests := []struct {
+		name           string
+		dataset        Dataset
+		threshold      float64
+		below          bool
+		expectedValues map[string]bool
+	}{
+		{
+			name:           "no values",
+			dataset:        make(map[string]float64),
+			threshold:      1,
+			below:          false,
+			expectedValues: make(map[string]bool),
+		},
+		{
+			name: "below and equal",
+			dataset: map[string]float64{
+				"a": 1,
+				"b": 2,
+				"c": 3,
+			},
+			threshold: 2,
+			below:     true,
+			expectedValues: map[string]bool{
+				"a": true,
+				"b": true,
+				"c": false,
+			},
+		},
+		{
+			name: "above and equal",
+			dataset: map[string]float64{
+				"a": 1,
+				"b": 2,
+				"c": 3,
+			},
+			threshold: 2,
+			below:     false,
+			expectedValues: map[string]bool{
+				"a": false,
+				"b": false,
+				"c": true,
+			},
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			values := test.dataset.GetThreshold(
+				test.threshold, test.below,
+			)
+
+			if !reflect.DeepEqual(test.expectedValues, values) {
+				t.Fatalf("expected: %v, got: %v",
+					test.expectedValues, values)
+			}
+		})
+	}
+}