From 860c2cdf43abb52d308059ca103f0fee620afe07 Mon Sep 17 00:00:00 2001
From: Ian Fijolek <ian@iamthefij.com>
Date: Wed, 12 May 2021 10:33:42 -0700
Subject: [PATCH] Add custom type to parse out seconds as int and durations as
 strings

---
 config.go       | 30 +++++++++++++++++++++++++++++-
 config_test.go  |  6 +++---
 main.go         |  2 +-
 monitor.go      |  8 ++++----
 monitor_test.go |  6 +++---
 5 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/config.go b/config.go
index 7d2d5b3..a389b2b 100644
--- a/config.go
+++ b/config.go
@@ -13,7 +13,7 @@ var errInvalidConfig = errors.New("Invalid configuration")
 
 // Config type is contains all provided user configuration
 type Config struct {
-	CheckInterval time.Duration `yaml:"check_interval"`
+	CheckInterval SecondsOrDuration `yaml:"check_interval"`
 	Monitors      []*Monitor
 	Alerts        map[string]*Alert
 }
@@ -52,6 +52,34 @@ func (cos *CommandOrShell) UnmarshalYAML(unmarshal func(interface{}) error) erro
 	return nil
 }
 
+// SecondsOrDuration wraps a duration value for parsing a duration or seconds from YAML
+// NOTE: This should be removed in favor of only parsing durations once compatibility is broken
+type SecondsOrDuration struct {
+	value time.Duration
+}
+
+// Value returns a duration value
+func (sod SecondsOrDuration) Value() time.Duration {
+	return sod.value
+}
+
+// UnmarshalYAML allows unmarshalling a duration value or seconds if an int was provided
+func (sod *SecondsOrDuration) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var seconds int64
+	err := unmarshal(&seconds)
+
+	if err == nil {
+		sod.value = time.Second * time.Duration(seconds)
+
+		return nil
+	}
+
+	// Error indicates that we don't have an int
+	err = unmarshal(&sod.value)
+
+	return err
+}
+
 // IsValid checks config validity and returns true if valid
 func (config Config) IsValid() (isValid bool) {
 	isValid = true
diff --git a/config_test.go b/config_test.go
index 8c5fab9..6945b40 100644
--- a/config_test.go
+++ b/config_test.go
@@ -52,15 +52,15 @@ func TestIntervalParsing(t *testing.T) {
 	oneMinute := time.Minute
 
 	// validate top level interval seconds represented as an int
-	if config.CheckInterval != oneSecond {
+	if config.CheckInterval.Value() != oneSecond {
 		t.Errorf("Incorrectly parsed int seconds. expected=%v actual=%v", oneSecond, config.CheckInterval)
 	}
 
-	if config.Monitors[0].CheckInterval != tenSeconds {
+	if config.Monitors[0].CheckInterval.Value() != tenSeconds {
 		t.Errorf("Incorrectly parsed seconds duration. expected=%v actual=%v", oneSecond, config.CheckInterval)
 	}
 
-	if config.Monitors[1].CheckInterval != oneMinute {
+	if config.Monitors[1].CheckInterval.Value() != oneMinute {
 		t.Errorf("Incorrectly parsed seconds duration. expected=%v actual=%v", oneSecond, config.CheckInterval)
 	}
 
diff --git a/main.go b/main.go
index 1eecd43..9121480 100644
--- a/main.go
+++ b/main.go
@@ -120,6 +120,6 @@ func main() {
 			panic(err)
 		}
 
-		time.Sleep(config.CheckInterval)
+		time.Sleep(config.CheckInterval.Value())
 	}
 }
diff --git a/monitor.go b/monitor.go
index 6aea6f3..a289897 100644
--- a/monitor.go
+++ b/monitor.go
@@ -11,9 +11,9 @@ import (
 // Monitor represents a particular periodic check of a command
 type Monitor struct { //nolint:maligned
 	// Config values
-	AlertAfter    int16         `yaml:"alert_after"`
-	AlertEvery    int16         `yaml:"alert_every"`
-	CheckInterval time.Duration `yaml:"check_interval"`
+	AlertAfter    int16             `yaml:"alert_after"`
+	AlertEvery    int16             `yaml:"alert_every"`
+	CheckInterval SecondsOrDuration `yaml:"check_interval"`
 	Name          string
 	AlertDown     []string `yaml:"alert_down"`
 	AlertUp       []string `yaml:"alert_up"`
@@ -45,7 +45,7 @@ func (monitor Monitor) ShouldCheck() bool {
 
 	sinceLastCheck := time.Since(monitor.lastCheck)
 
-	return sinceLastCheck >= monitor.CheckInterval
+	return sinceLastCheck >= monitor.CheckInterval.Value()
 }
 
 // Check will run the command configured by the Monitor and return a status
diff --git a/monitor_test.go b/monitor_test.go
index 5988b91..7d8f7ed 100644
--- a/monitor_test.go
+++ b/monitor_test.go
@@ -45,9 +45,9 @@ func TestMonitorShouldCheck(t *testing.T) {
 		name     string
 	}{
 		{Monitor{}, true, "Empty"},
-		{Monitor{lastCheck: timeNow, CheckInterval: time.Second * 15}, false, "Just checked"},
-		{Monitor{lastCheck: timeTenSecAgo, CheckInterval: time.Second * 15}, false, "-10s"},
-		{Monitor{lastCheck: timeTwentySecAgo, CheckInterval: time.Second * 15}, true, "-20s"},
+		{Monitor{lastCheck: timeNow, CheckInterval: SecondsOrDuration{time.Second * 15}}, false, "Just checked"},
+		{Monitor{lastCheck: timeTenSecAgo, CheckInterval: SecondsOrDuration{time.Second * 15}}, false, "-10s"},
+		{Monitor{lastCheck: timeTwentySecAgo, CheckInterval: SecondsOrDuration{time.Second * 15}}, true, "-20s"},
 	}
 
 	for _, c := range cases {