implement goyaml's unmarshaler

blackbox_suite_test.go

@@ -0,0 +1,13 @@
+package blackbox_test
+
+import (
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+
+	"testing"
+)
+
+func TestBlackbox(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Blackbox Suite")
+}

config.go

@@ -1,7 +1,7 @@
 package blackbox
 
 import (
-	"fmt"
+	"errors"
 	"io/ioutil"
 	"os"
 	"time"
@@ -12,15 +12,16 @@ import (
 
 type Duration time.Duration
 
-func (d *Duration) UnmarshalYAML(tag string, value interface{}) error {
-	if num, ok := value.(int64); ok {
+func (d *Duration) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var num int64
+	if err := unmarshal(&num); err == nil {
 		*d = Duration(num)
 		return nil
 	}
 
-	str, ok := value.(string)
-	if !ok {
-		return fmt.Errorf("invalid duration: %T (%#v)", value, value)
+	var str string
+	if err := unmarshal(&str); err != nil {
+		return errors.New("invalid duration; must be string or number")
 	}
 
 	duration, err := time.ParseDuration(str)

config_test.go

@@ -0,0 +1,31 @@
+package blackbox_test
+
+import (
+	"time"
+
+	. "github.com/concourse/blackbox"
+	"gopkg.in/yaml.v2"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Config", func() {
+	Describe("Duration", func() {
+		It("can be unmarshalled from YAML as a string", func() {
+			var duration Duration
+			err := yaml.Unmarshal([]byte("10s"), &duration)
+			Expect(err).ToNot(HaveOccurred())
+
+			Expect(time.Duration(duration)).To(Equal(10 * time.Second))
+		})
+
+		It("can be unmarshalled from YAML as an integer", func() {
+			var duration Duration
+			err := yaml.Unmarshal([]byte("10"), &duration)
+			Expect(err).ToNot(HaveOccurred())
+
+			Expect(time.Duration(duration)).To(Equal(10 * time.Nanosecond))
+		})
+	})
+})