From 0880b527384cea01dec5d63f0266b33427ddbccf Mon Sep 17 00:00:00 2001
From: Ivan Klymenchenko <developer@iillexial.me>
Date: Wed, 30 Dec 2020 20:41:35 +0200
Subject: [PATCH] Add context_test.go

---
 context.go      |  5 +++--
 context_test.go | 51 +++++++++++++++++++++++++++++++++++++++++++++++++
 main.go         |  6 +++---
 options_test.go |  9 ++++++++-
 4 files changed, 65 insertions(+), 6 deletions(-)
 create mode 100644 context_test.go

diff --git a/context.go b/context.go
index 6b6f4ee..c9faff2 100644
--- a/context.go
+++ b/context.go
@@ -6,8 +6,9 @@ type Context struct {
 	InsideTmuxSession bool
 }
 
-func CreateContext() *Context {
+func CreateContext() Context {
 	_, tmux := os.LookupEnv("TMUX")
+	os.Environ()
 	insideTmuxSession := os.Getenv("TERM") == "screen" || tmux
-	return &Context{insideTmuxSession}
+	return Context{insideTmuxSession}
 }
diff --git a/context_test.go b/context_test.go
new file mode 100644
index 0000000..02f6407
--- /dev/null
+++ b/context_test.go
@@ -0,0 +1,51 @@
+package main
+
+import (
+	"os"
+	"reflect"
+	"testing"
+)
+
+var environmentTestTable = []struct {
+	environment map[string]string
+	context     Context
+}{
+	{
+		map[string]string{},
+		Context{InsideTmuxSession: false},
+	},
+	{
+		map[string]string{
+			"TMUX": "",
+		},
+		Context{InsideTmuxSession: true},
+	},
+	{
+		map[string]string{
+			"TERM": "screen",
+		},
+		Context{InsideTmuxSession: true},
+	},
+	{
+		map[string]string{
+			"TERM": "xterm",
+			"TMUX": "",
+		},
+		Context{InsideTmuxSession: true},
+	},
+}
+
+func TestCreateContext(t *testing.T) {
+	os.Clearenv()
+	for _, v := range environmentTestTable {
+		for key, value := range v.environment {
+			os.Setenv(key, value)
+		}
+
+		context := CreateContext()
+
+		if !reflect.DeepEqual(v.context, context) {
+			t.Errorf("expected context %v, got %v", v.context, context)
+		}
+	}
+}
diff --git a/main.go b/main.go
index af5b6b5..46375d2 100644
--- a/main.go
+++ b/main.go
@@ -90,10 +90,10 @@ func main() {
 		} else {
 			fmt.Println("Starting new windows...")
 		}
-		err = smug.Start(*config, options, *context)
+		err = smug.Start(*config, options, context)
 		if err != nil {
 			fmt.Println("Oops, an error occurred! Rolling back...")
-			smug.Stop(*config, options, *context)
+			smug.Stop(*config, options, context)
 		}
 	case CommandStop:
 		if len(options.Windows) == 0 {
@@ -101,7 +101,7 @@ func main() {
 		} else {
 			fmt.Println("Killing windows...")
 		}
-		err = smug.Stop(*config, options, *context)
+		err = smug.Stop(*config, options, context)
 	}
 
 	if err != nil {
diff --git a/options_test.go b/options_test.go
index 0832dd7..abadf96 100644
--- a/options_test.go
+++ b/options_test.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"reflect"
 	"testing"
 
@@ -91,6 +92,12 @@ var usageTestTable = []struct {
 		ErrHelp,
 		1,
 	},
+	{
+		[]string{"start", "--test"},
+		Options{},
+		errors.New("unknown flag: --test"),
+		0,
+	},
 }
 
 func TestParseOptions(t *testing.T) {
@@ -116,7 +123,7 @@ func TestParseOptions(t *testing.T) {
 			t.Errorf("expected to get %d help calls, got %d", v.helpCalls, helpCalls)
 		}
 
-		if err != v.err {
+		if v.err != nil && err.Error() != v.err.Error() {
 			t.Errorf("expected to get error %v, got %v", v.err, err)
 		}