From a8d3dcef9a53e53b588e72908de07a691eae8286 Mon Sep 17 00:00:00 2001
From: cw9 <chaow929@gmail.com>
Date: Sat, 10 Jun 2017 14:24:36 -0400
Subject: [PATCH] YAML unmarshal aggregation types (#48)

---
 policy/aggregation_type.go      | 15 ++++++++++
 policy/aggregation_type_test.go | 53 ++++++++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/policy/aggregation_type.go b/policy/aggregation_type.go
index 31d627d..d6f5c83 100644
--- a/policy/aggregation_type.go
+++ b/policy/aggregation_type.go
@@ -217,6 +217,21 @@ func NewAggregationTypesFromSchema(input []schema.AggregationType) (AggregationT
 	return res, nil
 }
 
+// UnmarshalYAML unmarshals aggregation types from a string.
+func (aggTypes *AggregationTypes) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var str string
+	if err := unmarshal(&str); err != nil {
+		return err
+	}
+
+	parsed, err := ParseAggregationTypes(str)
+	if err != nil {
+		return err
+	}
+	*aggTypes = parsed
+	return nil
+}
+
 // IsDefault checks if the AggregationTypes is the default aggregation type.
 func (aggTypes AggregationTypes) IsDefault() bool {
 	return len(aggTypes) == 0
diff --git a/policy/aggregation_type_test.go b/policy/aggregation_type_test.go
index 4003408..348ad59 100644
--- a/policy/aggregation_type_test.go
+++ b/policy/aggregation_type_test.go
@@ -24,7 +24,9 @@ import (
 	"testing"
 
 	"github.com/m3db/m3x/pool"
+
 	"github.com/stretchr/testify/require"
+	yaml "gopkg.in/yaml.v2"
 )
 
 func TestAggregationTypeIsValid(t *testing.T) {
@@ -38,7 +40,56 @@ func TestAggregationTypesIsDefault(t *testing.T) {
 	require.False(t, AggregationTypes{Upper}.IsDefault())
 }
 
-func TestParseParseAggregationTypes(t *testing.T) {
+func TestAggregationTypesUnmarshalYAML(t *testing.T) {
+	inputs := []struct {
+		str         string
+		expected    AggregationTypes
+		expectedErr bool
+	}{
+		{
+			str:      "Lower",
+			expected: AggregationTypes{Lower},
+		},
+		{
+			str:      "Mean,Upper,P99,P9999",
+			expected: AggregationTypes{Mean, Upper, P99, P9999},
+		},
+		{
+			str:         "Lower,Upper,P99,P9999,P100",
+			expectedErr: true,
+		},
+		{
+			str:         "Lower,Upper,P99,P9999,P100",
+			expectedErr: true,
+		},
+		{
+			str:         ",Mean",
+			expectedErr: true,
+		},
+		{
+			str:         "Mean,",
+			expectedErr: true,
+		},
+		{
+			str:         ",Mean,",
+			expectedErr: true,
+		},
+	}
+	for _, input := range inputs {
+		var aggtypes AggregationTypes
+		err := yaml.Unmarshal([]byte(input.str), &aggtypes)
+
+		if input.expectedErr {
+			require.Error(t, err)
+			continue
+		}
+
+		require.NoError(t, err)
+		require.Equal(t, input.expected, aggtypes)
+	}
+}
+
+func TestParseAggregationTypes(t *testing.T) {
 	inputs := []struct {
 		str      string
 		expected AggregationTypes