From 496c70db0ef472ff2b04ae6486e4e050ced63c25 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Markus=20W=C3=BCstenberg?= <markus@maragu.dk>
Date: Thu, 24 Jun 2021 11:18:06 +0200
Subject: [PATCH] Add Load and MustLoad (#3)

---
 .github/workflows/go.yml |  4 +--
 README.md                |  2 +-
 env.go                   | 41 +++++++++++++++++++++++++++++++
 env_test.go              | 53 ++++++++++++++++++++++++++++++++++++++++
 testdata/env             |  2 ++
 testdata/invalid         |  1 +
 6 files changed, 100 insertions(+), 3 deletions(-)
 create mode 100644 testdata/env
 create mode 100644 testdata/invalid

diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
index 1f74973..550f78e 100644
--- a/.github/workflows/go.yml
+++ b/.github/workflows/go.yml
@@ -2,9 +2,9 @@ name: Go
 
 on:
   push:
-    branches: [ master ]
+    branches: [ main ]
   pull_request:
-    branches: [ master ]
+    branches: [ main ]
 
 jobs:
   build:
diff --git a/README.md b/README.md
index b6426cd..4f86a94 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,7 @@
 
 [![GoDoc](https://godoc.org/github.com/maragudk/env?status.svg)](https://godoc.org/github.com/maragudk/env)
 [![Go](https://github.com/maragudk/env/actions/workflows/go.yml/badge.svg)](https://github.com/maragudk/env/actions/workflows/go.yml)
-[![codecov](https://codecov.io/gh/maragudk/env/branch/master/graph/badge.svg)](https://codecov.io/gh/maragudk/env)
+[![codecov](https://codecov.io/gh/maragudk/env/branch/main/graph/badge.svg)](https://codecov.io/gh/maragudk/env)
 
 A small utility package to load different types of environment variables.
 
diff --git a/env.go b/env.go
index f67960a..5603215 100644
--- a/env.go
+++ b/env.go
@@ -2,8 +2,11 @@
 package env
 
 import (
+	"bufio"
+	"fmt"
 	"os"
 	"strconv"
+	"strings"
 	"time"
 )
 
@@ -54,3 +57,41 @@ func GetDurationOrDefault(name string, defaultV time.Duration) time.Duration {
 	}
 	return vAsDuration
 }
+
+// Load environment variables from environment files. Defaults to loading from .env.
+func Load(paths ...string) error {
+	if len(paths) == 0 {
+		paths = []string{".env"}
+	}
+	for _, path := range paths {
+		file, err := os.Open(path)
+		if err != nil {
+			return fmt.Errorf("error reading %v: %w", path, err)
+		}
+		s := bufio.NewScanner(file)
+		var i int
+		for s.Scan() {
+			i++
+			line := s.Text()
+			parts := strings.Split(line, "=")
+			if len(parts) < 2 {
+				return fmt.Errorf("missing equal sign on line %v in %v", i, path)
+			}
+			if err := os.Setenv(parts[0], parts[1]); err != nil {
+				return fmt.Errorf("error setting environment variable for line %v in %v: %w", i, path, err)
+			}
+		}
+		if err := s.Err(); err != nil {
+			return fmt.Errorf("error scanning %v: %w", path, err)
+		}
+		_ = file.Close()
+	}
+	return nil
+}
+
+// MustLoad calls Load and panics on errors.
+func MustLoad(paths ...string) {
+	if err := Load(paths...); err != nil {
+		panic(err)
+	}
+}
diff --git a/env_test.go b/env_test.go
index ae88989..5cbc6c4 100644
--- a/env_test.go
+++ b/env_test.go
@@ -90,6 +90,51 @@ func TestGetDurationOrDefault(t *testing.T) {
 	})
 }
 
+func TestLoad(t *testing.T) {
+	t.Run("loads an environment file", func(t *testing.T) {
+		is := is.New(t)
+		defer unsetenv("hat", "hats")
+		err := env.Load("testdata/env")
+		is.NoErr(err)
+		hat := env.GetStringOrDefault("hat", "regular")
+		is.Equal("party", hat)
+		hats := env.GetIntOrDefault("hats", 1)
+		is.Equal(2, hats)
+	})
+
+	t.Run("errors on bad file", func(t *testing.T) {
+		is := is.New(t)
+		err := env.Load("testdata/invalid")
+		is.True(err != nil)
+		is.Equal("missing equal sign on line 1 in testdata/invalid", err.Error())
+	})
+}
+
+func TestMustLoad(t *testing.T) {
+	t.Run("loads an environment file", func(t *testing.T) {
+		is := is.New(t)
+		defer unsetenv("hat", "hats")
+		env.MustLoad("testdata/env")
+		hat := env.GetStringOrDefault("hat", "regular")
+		is.Equal("party", hat)
+		hats := env.GetIntOrDefault("hats", 1)
+		is.Equal(2, hats)
+	})
+
+	t.Run("panics on no such file", func(t *testing.T) {
+		is := is.New(t)
+		recovered := false
+		defer func() {
+			if err := recover(); err != nil {
+				recovered = true
+			}
+			is.True(recovered)
+		}()
+		env.MustLoad()
+
+	})
+}
+
 func setenv(k, v string) func() {
 	if err := os.Setenv(k, v); err != nil {
 		panic(err)
@@ -100,3 +145,11 @@ func setenv(k, v string) func() {
 		}
 	}
 }
+
+func unsetenv(ks ...string) {
+	for _, k := range ks {
+		if err := os.Unsetenv(k); err != nil {
+			panic(err)
+		}
+	}
+}
diff --git a/testdata/env b/testdata/env
new file mode 100644
index 0000000..63ae506
--- /dev/null
+++ b/testdata/env
@@ -0,0 +1,2 @@
+hat=party
+hats=2
diff --git a/testdata/invalid b/testdata/invalid
new file mode 100644
index 0000000..d58eb7b
--- /dev/null
+++ b/testdata/invalid
@@ -0,0 +1 @@
+notvalid