Merge pull request #28 from kluctl/feat-custom-tmp-dir

fix: Fix custom python interpreter extract dir and flaky tests
kluctl · Jan 8, 2024 · 8839259 · 8839259
2 parents 4b0aaea + 90d67f0
commit 8839259
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 38 deletions.
diff --git a/jinja2.go b/jinja2.go
@@ -69,17 +69,17 @@ func NewJinja2(name string, parallelism int, opts ...Jinja2Opt) (*Jinja2, error)
 	}
 	tmpDir = filepath.Join(tmpDir, name)
 
-	j2.ep, err = python.NewEmbeddedPython(name)
+	j2.ep, err = python.NewEmbeddedPythonWithTmpDir(tmpDir+"-python", true)
 	if err != nil {
 		return nil, err
 	}
-	j2.jinja2Lib, err = embed_util.NewEmbeddedFilesWithTmpDir(data.Data, tmpDir+"-python", true)
+	j2.jinja2Lib, err = embed_util.NewEmbeddedFilesWithTmpDir(data.Data, tmpDir+"-jinja2-lib", true)
 	if err != nil {
 		return nil, err
 	}
 	j2.ep.AddPythonPath(j2.jinja2Lib.GetExtractedPath())
 
-	j2.rendererSrc, err = embed_util.NewEmbeddedFilesWithTmpDir(python_src.RendererSource, tmpDir+"-renderer", true)
+	j2.rendererSrc, err = embed_util.NewEmbeddedFilesWithTmpDir(python_src.RendererSource, tmpDir+"-jinja2-renderer", true)
 	if err != nil {
 		return nil, err
 	}

diff --git a/jinja2_renderer.go b/jinja2_renderer.go
@@ -58,14 +58,21 @@ func newPythonJinja2Renderer(j2 *Jinja2) (*pythonJinja2Renderer, error) {
 
 	j2r.stdoutReader = bufio.NewReader(j2r.stdout)
 
+	_, err = j2r.runCmd(&jinja2Cmd{
+		Cmd: "init",
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to initialize renderer: %w", err)
+	}
+
 	isOk = true
 
 	return j2r, nil
 }
 
 func (j *pythonJinja2Renderer) Close() {
 	if j.stdin != nil {
-		args := jinja2Args{Cmd: "exit"}
+		args := jinja2Cmd{Cmd: "exit"}
 		_ = json.NewEncoder(j.stdin).Encode(args)
 
 		_ = j.stdin.Close()
@@ -133,20 +140,24 @@ func isMaybeTemplate(template string, isString bool) (bool, *string, error) {
 	return true, nil, nil
 }
 
-type jinja2Args struct {
+type jinja2Cmd struct {
 	Cmd       string   `json:"cmd"`
 	Templates []string `json:"templates"`
 
 	Opts *jinja2Options `json:"opts"`
 }
 
-type jinja2Result struct {
+type jinja2TemplateResult struct {
 	Result *string `json:"result,omitempty"`
 	Error  *string `json:"error,omitempty"`
 }
 
+type jinja2CmdResult struct {
+	TemplateResults []jinja2TemplateResult `json:"templateResults,omitempty"`
+}
+
 func (j *pythonJinja2Renderer) renderHelper(jobs []*RenderJob, isString bool, opts []Jinja2Opt) error {
-	var jargs jinja2Args
+	var jargs jinja2Cmd
 	if isString {
 		jargs.Cmd = "render-strings"
 	} else {
@@ -191,14 +202,34 @@ func (j *pythonJinja2Renderer) renderHelper(jobs []*RenderJob, isString bool, op
 		return nil
 	}
 
+	cmdResult, err := j.runCmd(&jargs)
+	if err != nil {
+		return err
+	}
+
+	for i, item := range cmdResult.TemplateResults {
+		if item.Result != nil {
+			processedJobs[i].Result = item.Result
+		} else {
+			if item.Error == nil {
+				return fmt.Errorf("missing result and error from item at index %d", i)
+			}
+			processedJobs[i].Error = &Jinja2Error{*item.Error}
+		}
+	}
+
+	return nil
+}
+
+func (j *pythonJinja2Renderer) runCmd(jargs *jinja2Cmd) (*jinja2CmdResult, error) {
 	b, err := json.Marshal(jargs)
 	if err != nil {
 		j.Close()
-		return err
+		return nil, err
 	}
 	b = append(b, '\n')
 
-	if jargs.Opts.traceJsonSend != nil {
+	if jargs.Opts != nil && jargs.Opts.traceJsonSend != nil {
 		var m map[string]any
 		_ = json.Unmarshal(b, &m)
 		jargs.Opts.traceJsonSend(m)
@@ -207,43 +238,32 @@ func (j *pythonJinja2Renderer) renderHelper(jobs []*RenderJob, isString bool, op
 	_, err = j.stdin.Write(b)
 	if err != nil {
 		j.Close()
-		return err
+		return nil, err
 	}
 
 	line := bytes.NewBuffer(nil)
 	for true {
 		l, p, err := j.stdoutReader.ReadLine()
 		if err != nil {
-			return err
+			return nil, err
 		}
 		line.Write(l)
 		if !p {
 			break
 		}
 	}
 
-	if jargs.Opts.traceJsonReceive != nil {
+	if jargs.Opts != nil && jargs.Opts.traceJsonReceive != nil {
 		var m map[string]any
 		_ = json.Unmarshal(line.Bytes(), &m)
 		jargs.Opts.traceJsonReceive(m)
 	}
 
-	var result []jinja2Result
+	var result jinja2CmdResult
 	err = json.Unmarshal(line.Bytes(), &result)
 	if err != nil {
-		return err
-	}
-
-	for i, item := range result {
-		if item.Result != nil {
-			processedJobs[i].Result = item.Result
-		} else {
-			if item.Error == nil {
-				return fmt.Errorf("missing result and error from item at index %d", i)
-			}
-			processedJobs[i].Error = &Jinja2Error{*item.Error}
-		}
+		return nil, err
 	}
 
-	return nil
+	return &result, nil
 }
diff --git a/jinja2_test.go b/jinja2_test.go
@@ -459,7 +459,10 @@ func TestCustomTmpDir(t *testing.T) {
 	m, err := filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-python-*"))
 	assert.NoError(t, err)
 	assert.Len(t, m, 2) // dir + .lock
-	m, err = filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-renderer-*"))
+	m, err = filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-jinja2-lib-*"))
+	assert.NoError(t, err)
+	assert.Len(t, m, 2) // dir + .lock
+	m, err = filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-jinja2-renderer-*"))
 	assert.NoError(t, err)
 	assert.Len(t, m, 2) // dir + .lock
 
@@ -473,13 +476,19 @@ func TestCustomTmpDir(t *testing.T) {
 	m, err = filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-python-*"))
 	assert.NoError(t, err)
 	assert.Len(t, m, 0)
-	m, err = filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-renderer-*"))
+	m, err = filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-jinja2-lib-*"))
+	assert.NoError(t, err)
+	assert.Len(t, m, 0)
+	m, err = filepath.Glob(filepath.Join(os.TempDir(), "go-jinja2-embedded", j2.name+"-jinja2-renderer-*"))
 	assert.NoError(t, err)
 	assert.Len(t, m, 0)
 	m, err = filepath.Glob(filepath.Join(tmpDir, j2.name+"-python-*"))
 	assert.NoError(t, err)
 	assert.Len(t, m, 2) // dir + .lock
-	m, err = filepath.Glob(filepath.Join(tmpDir, j2.name+"-renderer-*"))
+	m, err = filepath.Glob(filepath.Join(tmpDir, j2.name+"-jinja2-lib-*"))
+	assert.NoError(t, err)
+	assert.Len(t, m, 2) // dir + .lock
+	m, err = filepath.Glob(filepath.Join(tmpDir, j2.name+"-jinja2-renderer-*"))
 	assert.NoError(t, err)
 	assert.Len(t, m, 2) // dir + .lock
 

diff --git a/python_src/main.py b/python_src/main.py
@@ -6,19 +6,26 @@
 
 def main():
     while True:
-        args = sys.stdin.readline()
-        if not args:
+        cmd = sys.stdin.readline()
+        if not cmd:
             break
-        args = json.loads(args)
-        opts = args["opts"]
+        cmd = json.loads(cmd)
+        opts = cmd["opts"]
 
         r = Jinja2Renderer(opts)
 
-        if args["cmd"] == "render-strings":
-            result = r.RenderStrings(args["templates"])
-        elif args["cmd"] == "render-files":
-            result = r.RenderFiles(args["templates"])
-        elif args["cmd"] == "exit":
+        result = {}
+        if cmd["cmd"] == "init":
+            pass
+        elif cmd["cmd"] == "render-strings":
+            result = {
+                "templateResults": r.RenderStrings(cmd["templates"])
+            }
+        elif cmd["cmd"] == "render-files":
+            result = {
+                "templateResults": r.RenderFiles(cmd["templates"])
+            }
+        elif cmd["cmd"] == "exit":
             break
         else:
             raise Exception("invalid cmd")