From 10ba98c6e2001440d3436fc0ef0d4420b2c07c37 Mon Sep 17 00:00:00 2001
From: lyyyuna <lyyyuna@outlook.com>
Date: Thu, 22 Jul 2021 22:25:55 +0800
Subject: [PATCH] modify cmd interface

---
 tests/e2e/cmd.go      | 11 +++++++++--
 tests/e2e/e2e_test.go |  6 +++---
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/tests/e2e/cmd.go b/tests/e2e/cmd.go
index 2ade2f0..49590b3 100644
--- a/tests/e2e/cmd.go
+++ b/tests/e2e/cmd.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"os"
 	"os/exec"
 	"time"
 )
@@ -18,10 +19,15 @@ type LongRunCmd struct {
 	done      bool
 }
 
-func NewLongRunCmd(dir string, args []string) *LongRunCmd {
+// NewLongRunCmd defines a command which will be run forever
+//
+// You can specify the whole command arg list, the directory where the command
+// will be run in, and the env list.
+func NewLongRunCmd(args []string, dir string, envs []string) *LongRunCmd {
 	ctx, cancel := context.WithCancel(context.Background())
 	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
 	cmd.Dir = dir
+	cmd.Env = append(os.Environ(), envs...)
 
 	var stderrBuf bytes.Buffer
 	var stdoutBuf bytes.Buffer
@@ -71,12 +77,13 @@ func (l *LongRunCmd) GetStdoutStdErr() (string, string) {
 }
 
 // RunShortRunCmd defines a cmd which run and exits immediately
-func RunShortRunCmd(dir string, args []string) (string, error) {
+func RunShortRunCmd(args []string, dir string, envs []string) (string, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*20)
 	defer cancel()
 
 	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
 	cmd.Dir = dir
+	cmd.Env = append(os.Environ(), envs...)
 
 	output, err := cmd.CombinedOutput()
 	return string(output), err
diff --git a/tests/e2e/e2e_test.go b/tests/e2e/e2e_test.go
index db9108f..522d493 100644
--- a/tests/e2e/e2e_test.go
+++ b/tests/e2e/e2e_test.go
@@ -20,7 +20,7 @@ var _ = Describe("1 [基础测试]", func() {
 			Expect(err).To(BeNil(), "找不到 sample")
 
 			By("使用 goc build 命令编译")
-			_, err = RunShortRunCmd(dir, []string{"goc", "build", "."})
+			_, err = RunShortRunCmd([]string{"goc", "build", "."}, dir, nil)
 			Expect(err).To(BeNil(), "goc build 运行错误")
 		})
 	})
@@ -31,12 +31,12 @@ var _ = Describe("1 [基础测试]", func() {
 			Expect(err).To(BeNil(), "找不到 sample")
 
 			By("启动 goc server")
-			lc := NewLongRunCmd(dir, []string{"goc", "server", "."})
+			lc := NewLongRunCmd([]string{"goc", "server", "."}, dir, nil)
 			lc.Run()
 			defer lc.Stop()
 
 			By("使用 goc list 获取服务列表")
-			output, err := RunShortRunCmd(dir, []string{"goc", "list"})
+			output, err := RunShortRunCmd([]string{"goc", "list"}, dir, nil)
 			Expect(err).To(BeNil(), "goc list 运行错误")
 			Expect(output).To(ContainSubstring("REMOTEIP"), "goc list 输出应该包含 REMOTEIP")
 		})