From 809df3fa66be7ca74005f668b8693eda336d50df Mon Sep 17 00:00:00 2001
From: Zygmunt Krynicki <zygmunt.krynicki@canonical.com>
Date: Mon, 2 May 2016 16:30:08 +0200
Subject: [PATCH] timeout: add YAML unmarshaler

Signed-off-by: Zygmunt Krynicki <zygmunt.krynicki@canonical.com>
---
 timeout/timeout.go      | 13 +++++++++++++
 timeout/timeout_test.go | 11 +++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/timeout/timeout.go b/timeout/timeout.go
index 0c812064623..265ae64c454 100644
--- a/timeout/timeout.go
+++ b/timeout/timeout.go
@@ -52,6 +52,19 @@ func (t *Timeout) UnmarshalJSON(buf []byte) error {
 	return nil
 }
 
+func (t *Timeout) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var str string
+	if err := unmarshal(&str); err != nil {
+		return err
+	}
+	dur, err := time.ParseDuration(str)
+	if err != nil {
+		return err
+	}
+	*t = Timeout(dur)
+	return nil
+}
+
 // String returns a string representing the duration
 func (t Timeout) String() string {
 	return time.Duration(t).String()
diff --git a/timeout/timeout_test.go b/timeout/timeout_test.go
index a1477a55d25..e58a6e0fc96 100644
--- a/timeout/timeout_test.go
+++ b/timeout/timeout_test.go
@@ -25,6 +25,7 @@ import (
 	"time"
 
 	. "gopkg.in/check.v1"
+	"gopkg.in/yaml.v2"
 )
 
 // Hook up check.v1 into the "go test" runner
@@ -42,7 +43,7 @@ func (s *TimeoutTestSuite) TestTimeoutMarshal(c *C) {
 }
 
 type testT struct {
-	T Timeout
+	T Timeout `yaml:"T"`
 }
 
 func (s *TimeoutTestSuite) TestTimeoutMarshalIndirect(c *C) {
@@ -51,8 +52,14 @@ func (s *TimeoutTestSuite) TestTimeoutMarshalIndirect(c *C) {
 	c.Check(string(bs), Equals, `{"T":"30s"}`)
 }
 
-func (s *TimeoutTestSuite) TestTimeoutUnmarshal(c *C) {
+func (s *TimeoutTestSuite) TestTimeoutUnmarshalJSON(c *C) {
 	var t testT
 	c.Assert(json.Unmarshal([]byte(`{"T": "17ms"}`), &t), IsNil)
 	c.Check(t, DeepEquals, testT{T: Timeout(17 * time.Millisecond)})
 }
+
+func (s *TimeoutTestSuite) TestTimeoutUnmarshalYAML(c *C) {
+	var t testT
+	c.Assert(yaml.Unmarshal([]byte(`T: 17ms`), &t), IsNil)
+	c.Check(t, DeepEquals, testT{T: Timeout(17 * time.Millisecond)})
+}