From c155a91098e3c16721433130c82c3525abe4a450 Mon Sep 17 00:00:00 2001
From: Zachary Wasserman <zachwass2000@gmail.com>
Date: Tue, 22 Jan 2019 18:30:48 -0800
Subject: [PATCH] Add env.Int function for retrieving integer environment
 variables (#43)

---
 env/env.go      | 20 ++++++++++++++++++--
 env/env_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/env/env.go b/env/env.go
index 9a3047c..c7d2a31 100644
--- a/env/env.go
+++ b/env/env.go
@@ -17,6 +17,7 @@ package env
 import (
 	"fmt"
 	"os"
+	"strconv"
 	"time"
 )
 
@@ -29,8 +30,23 @@ func String(key, def string) string {
 	return def
 }
 
-// Bool returns the environment variable value specified by the key parameter,
-// otherwise returning a default value if set.
+// Int returns the environment variable value specified by the key parameter,
+// parsed as an integer. If the environment variable is not set, the default
+// value is returned. If parsing the integer fails, Int will exit the program.
+func Int(key string, def int) int {
+	if env, ok := os.LookupEnv(key); ok {
+		i, err := strconv.Atoi(env)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "env: parse int from flag: %s\n", err)
+			os.Exit(1)
+		}
+		return i
+	}
+	return def
+}
+
+// Bool returns the environment variable value specified by the key parameter
+// (parsed as a boolean), otherwise returning a default value if set.
 func Bool(key string, def bool) bool {
 	env, ok := os.LookupEnv(key)
 	if !ok {
diff --git a/env/env_test.go b/env/env_test.go
index 70d5bf9..0711925 100644
--- a/env/env_test.go
+++ b/env/env_test.go
@@ -105,3 +105,34 @@ func TestBool(t *testing.T) {
 		t.Errorf("have %v, want %v", have, want)
 	}
 }
+
+func TestInt(t *testing.T) {
+	var tests = []struct {
+		env   string
+		value int
+	}{
+		{env: "1337", value: 1337},
+		{env: "1", value: 1},
+		{env: "-34", value: -34},
+		{env: "0", value: 0},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.env, func(t *testing.T) {
+			key := "TEST_INT"
+			if err := os.Setenv(key, tt.env); err != nil {
+				t.Fatalf("failed to set env var %s for test: %s\n", key, err)
+			}
+
+			if have, want := Int(key, 10), tt.value; have != want {
+				t.Errorf("have %v, want %v", have, want)
+			}
+		})
+	}
+
+	// test default value
+	def := 11
+	if have, want := Int("TEST_DEFAULT", def), def; have != want {
+		t.Errorf("have %v, want %v", have, want)
+	}
+}