From fbc9f8e6409dd0122d49ccf8a8e7696c80de442c Mon Sep 17 00:00:00 2001
From: Faissal Elamraoui <amr.faissal@gmail.com>
Date: Fri, 13 Jan 2017 17:10:54 +0100
Subject: [PATCH] Package env for working with env variables inside Beego

The package env makes it trivial to work with environment variables.
It allows getting values with defaults as a fallback.
New ENV variables can be set safely on the current process environment.
---
 env/env.go      | 89 +++++++++++++++++++++++++++++++++++++++++++++++++
 env/env_test.go | 75 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 164 insertions(+)
 create mode 100644 env/env.go
 create mode 100644 env/env_test.go

diff --git a/env/env.go b/env/env.go
new file mode 100644
index 00000000..014501f5
--- /dev/null
+++ b/env/env.go
@@ -0,0 +1,89 @@
+// Copyright 2014 beego Author. All Rights Reserved.
+// Copyright 2017 Faissal Elamraoui. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package env
+
+import (
+	"fmt"
+	"os"
+	"strings"
+	"sync"
+)
+
+var env struct {
+	data map[string]string
+	lock *sync.RWMutex
+}
+
+func init() {
+	env.data = make(map[string]string)
+	env.lock = &sync.RWMutex{}
+	for _, e := range os.Environ() {
+		splits := strings.Split(e, "=")
+		env.data[splits[0]] = os.Getenv(splits[0])
+	}
+}
+
+// Get returns a value by key.
+// If the key does not exist, the default value will be returned.
+func Get(key string, defVal string) string {
+	env.lock.RLock()
+	defer env.lock.RUnlock()
+
+	if val, ok := env.data[key]; ok {
+		return val
+	}
+	return defVal
+}
+
+// MustGet returns a value by key.
+// If the key does not exist, it will return an error.
+func MustGet(key string) (string, error) {
+	env.lock.RLock()
+	defer env.lock.RUnlock()
+
+	if val, ok := env.data[key]; ok {
+		return val, nil
+	}
+	return "", fmt.Errorf("no env variable with %s", key)
+}
+
+// Set sets a value in the ENV copy.
+// This does not affect the child process environment.
+func Set(key string, value string) {
+	env.lock.Lock()
+	defer env.lock.Unlock()
+	env.data[key] = value
+}
+
+// MustSet sets a value in the ENV copy and the child process environment.
+// It returns an error in case the set operation failed.
+func MustSet(key string, value string) error {
+	env.lock.Lock()
+	defer env.lock.Unlock()
+
+	err := os.Setenv(key, value)
+	if err != nil {
+		return err
+	}
+	env.data[key] = value
+	return nil
+}
+
+// GetAll returns all keys/values in the current child process environment.
+func GetAll() map[string]string {
+	env.lock.RLock()
+	defer env.lock.RUnlock()
+	return env.data
+}
diff --git a/env/env_test.go b/env/env_test.go
new file mode 100644
index 00000000..3f1d4dba
--- /dev/null
+++ b/env/env_test.go
@@ -0,0 +1,75 @@
+// Copyright 2014 beego Author. All Rights Reserved.
+// Copyright 2017 Faissal Elamraoui. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package env
+
+import (
+	"os"
+	"testing"
+)
+
+func TestEnvGet(t *testing.T) {
+	gopath := Get("GOPATH", "")
+	if gopath != os.Getenv("GOPATH") {
+		t.Error("expected GOPATH not empty.")
+	}
+
+	noExistVar := Get("NOEXISTVAR", "foo")
+	if noExistVar != "foo" {
+		t.Errorf("expected NOEXISTVAR to equal foo, got %s.", noExistVar)
+	}
+}
+
+func TestEnvMustGet(t *testing.T) {
+	gopath, err := MustGet("GOPATH")
+	if err != nil {
+		t.Error(err)
+	}
+
+	if gopath != os.Getenv("GOPATH") {
+		t.Errorf("expected GOPATH to be the same, got %s.", gopath)
+	}
+
+	_, err = MustGet("NOEXISTVAR")
+	if err == nil {
+		t.Error("expected error to be non-nil")
+	}
+}
+
+func TestEnvSet(t *testing.T) {
+	Set("MYVAR", "foo")
+	myVar := Get("MYVAR", "bar")
+	if myVar != "foo" {
+		t.Errorf("expected MYVAR to equal foo, got %s.", myVar)
+	}
+}
+
+func TestEnvMustSet(t *testing.T) {
+	err := MustSet("FOO", "bar")
+	if err != nil {
+		t.Error(err)
+	}
+
+	fooVar := os.Getenv("FOO")
+	if fooVar != "bar" {
+		t.Errorf("expected FOO variable to equal bar, got %s.", fooVar)
+	}
+}
+
+func TestEnvGetAll(t *testing.T) {
+	envMap := GetAll()
+	if len(envMap) == 0 {
+		t.Error("expected environment not empty.")
+	}
+}