From a87c8336aaac26a85bde96c2c2ea0c4d0a7a56c2 Mon Sep 17 00:00:00 2001
From: "Serge A. Zaitsev" <zaitsev.serge@gmail.com>
Date: Fri, 5 Dec 2014 20:09:03 +0200
Subject: [PATCH] render uses strings, not bytes; added helpers for env and run

---
 zs.go      | 75 +++++++++++++++++++++++++++++++-----------------------
 zs_test.go | 57 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 99 insertions(+), 33 deletions(-)

diff --git a/zs.go b/zs.go
index 12cddbd..697d778 100644
--- a/zs.go
+++ b/zs.go
@@ -46,32 +46,30 @@ func md(s string) (map[string]string, string) {
 }
 
 func render(s string, vars map[string]string, eval EvalFn) (string, error) {
-	b := []byte(s)
-	delim_open := []byte("{{")
-	delim_close := []byte("}}")
+	delim_open := "{{"
+	delim_close := "}}"
 
 	out := bytes.NewBuffer(nil)
 	for {
-		if from := bytes.Index(b, delim_open); from == -1 {
-			out.Write(b)
+		if from := strings.Index(s, delim_open); from == -1 {
+			out.WriteString(s)
 			return out.String(), nil
 		} else {
-			to := bytes.Index(b, delim_close)
-			if to == -1 {
+			if to := strings.Index(s, delim_close); to == -1 {
 				return "", fmt.Errorf("Close delim not found")
 			} else {
-				out.Write(b[:from])
-				cmd := b[from+len(delim_open) : to]
-				b = b[to+len(delim_close):]
-				m := strings.Fields(string(cmd))
+				out.WriteString(s[:from])
+				cmd := s[from+len(delim_open) : to]
+				s = s[to+len(delim_close):]
+				m := strings.Fields(cmd)
 				if len(m) == 1 {
 					if v, ok := vars[m[0]]; ok {
-						out.Write([]byte(v))
+						out.WriteString(v)
 						continue
 					}
 				}
 				if res, err := eval(m, vars); err == nil {
-					out.Write([]byte(res))
+					out.WriteString(res)
 				} else {
 					log.Println(err) // silent
 				}
@@ -80,29 +78,45 @@ func render(s string, vars map[string]string, eval EvalFn) (string, error) {
 	}
 }
 
-func eval(cmd []string, vars map[string]string) (string, error) {
-	var outbuf, errbuf bytes.Buffer
-	c := exec.Command(path.Join(ZSDIR, cmd[0]), cmd[1:]...)
+func env(vars map[string]string) []string {
 	env := []string{"ZS=" + os.Args[0]}
+	env = append(env, os.Environ()...)
 	for k, v := range vars {
 		env = append(env, "ZS_"+strings.ToUpper(k)+"="+v)
 	}
-	c.Env = append(c.Env, env...)
-	c.Stdout = &outbuf
+	return env
+}
+
+func run(cmd string, args []string, vars map[string]string, output io.Writer) error {
+	var errbuf bytes.Buffer
+	c := exec.Command(cmd, args...)
+	c.Env = env(vars)
+	c.Stdout = output
 	c.Stderr = &errbuf
-	if err := c.Run(); err != nil {
+
+	err := c.Run()
+
+	if errbuf.Len() > 0 {
+		log.Println(errbuf.String())
+	}
+
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func eval(cmd []string, vars map[string]string) (string, error) {
+	outbuf := bytes.NewBuffer(nil)
+	err := run(path.Join(ZSDIR, cmd[0]), cmd[1:], vars, outbuf)
+	if err != nil {
 		log.Println(err)
-		c := exec.Command(path.Join(cmd[0]), cmd[1:]...)
-		c.Env = append(c.Env, env...)
-		c.Stdout = &outbuf
-		c.Stderr = &errbuf
-		if err := c.Run(); err != nil {
+		outbuf = bytes.NewBuffer(nil)
+		err := run(cmd[0], cmd[1:], vars, outbuf)
+		if err != nil {
 			return "", err
 		}
 	}
-	if errbuf.Len() > 0 {
-		log.Println(errbuf.String())
-	}
 	return outbuf.String(), nil
 }
 
@@ -247,11 +261,8 @@ func main() {
 			log.Println(err)
 		}
 	default:
-		cmd := exec.Command(path.Join(ZSDIR, cmd), args...)
-		cmd.Env = append(cmd.Env, "ZS="+os.Args[0])
-		cmd.Stdout = os.Stdout
-		cmd.Stderr = os.Stderr
-		if err := cmd.Run(); err != nil {
+		err := run(path.Join(ZSDIR, cmd), args, map[string]string{}, os.Stdout)
+		if err != nil {
 			log.Println(err)
 		}
 	}
diff --git a/zs_test.go b/zs_test.go
index 1d3af9e..b4ee3b1 100644
--- a/zs_test.go
+++ b/zs_test.go
@@ -1,6 +1,13 @@
 package main
 
-import "testing"
+import (
+	"bytes"
+	"fmt"
+	"log"
+	"os"
+	"strings"
+	"testing"
+)
 
 func TestSplit2(t *testing.T) {
 	if a, b := split2("a:b", ":"); a != "a" || b != "b" {
@@ -75,4 +82,52 @@ func TestRender(t *testing.T) {
 	if s, err := render("{{greet}} x{{foo}}z", vars, eval); err != nil || s != "hello xbarz" {
 		t.Error()
 	}
+	// Test error case
+	if s, err := render("a {{greet text ", vars, eval); err == nil || len(s) != 0 {
+		t.Error()
+	}
+}
+
+func TestEnv(t *testing.T) {
+	e := env(map[string]string{"foo": "bar", "baz": "hello world"})
+	mustHave := []string{"ZS=" + os.Args[0], "ZS_FOO=bar", "ZS_BAZ=hello world", "PATH="}
+	for _, s := range mustHave {
+		found := false
+		for _, v := range e {
+			if strings.HasPrefix(v, s) {
+				found = true
+				break
+			}
+		}
+		if !found {
+			t.Error("Missing", s)
+		}
+	}
+}
+
+func TestRun(t *testing.T) {
+	out := bytes.NewBuffer(nil)
+	err := run("some_unbelievable_command_name", []string{}, map[string]string{}, out)
+	if err == nil {
+		t.Error()
+	}
+
+	out = bytes.NewBuffer(nil)
+	err = run(os.Args[0], []string{"-test.run=TestHelperProcess"},
+		map[string]string{"helper": "1", "out": "foo", "err": "bar"}, out)
+	if err != nil {
+		t.Error(err)
+	}
+	if out.String() != "foo\n" {
+		t.Error(out.String())
+	}
+}
+
+func TestHelperProcess(*testing.T) {
+	if os.Getenv("ZS_HELPER") != "1" {
+		return
+	}
+	defer os.Exit(0)                 // TODO check exit code
+	log.Println(os.Getenv("ZS_ERR")) // stderr
+	fmt.Println(os.Getenv("ZS_OUT")) // stdout
 }
-- 
2.43.0