From 8a3fcc0302000057504c0a964355a1c7974e8551 Mon Sep 17 00:00:00 2001
From: Krishnan Anantheswaran <krishnan@the-krishnans.com>
Date: Wed, 3 Feb 2021 13:55:13 -0800
Subject: [PATCH] Allow reset of ext and TLA vars for VM reuse (#509)

Allow reset of ext and TLA vars for VM reuse

Currently ext and TLA vars can be set but not reset. This makes
re-using VMs that have different variables impossible.

Add VM methods to be able to reset ext and TLA vars
---
 jsonnet_test.go | 37 +++++++++++++++++++++++++++++++++++++
 vm.go           | 11 +++++++++++
 2 files changed, 48 insertions(+)

diff --git a/jsonnet_test.go b/jsonnet_test.go
index b5b5d97..fbdefd3 100644
--- a/jsonnet_test.go
+++ b/jsonnet_test.go
@@ -3,6 +3,7 @@ package jsonnet
 import (
 	"bytes"
 	"reflect"
+	"strings"
 	"testing"
 	"unicode/utf8"
 
@@ -228,3 +229,39 @@ func TestContents(t *testing.T) {
 		t.Errorf("Contents should distinguish between different instances even if they have the same data inside")
 	}
 }
+
+func TestExtReset(t *testing.T) {
+	vm := MakeVM()
+	vm.ExtVar("fooString", "bar")
+	vm.ExtCode("fooCode", "true")
+	_, err := vm.EvaluateAnonymousSnippet("test.jsonnet", `{ str: std.extVar('fooString'), code: std.extVar('fooCode') }`)
+	if err != nil {
+		t.Fatalf("unexpected error %v", err)
+	}
+	vm.ExtReset()
+	_, err = vm.EvaluateAnonymousSnippet("test.jsonnet", `{ str: std.extVat('fooString'), code: std.extVar('fooCode') }`)
+	if err == nil {
+		t.Fatalf("expected error, got nil")
+	}
+	if !strings.Contains(err.Error(), "Undefined external variable") {
+		t.Errorf("unexpected error %v", err)
+	}
+}
+
+func TestTLAReset(t *testing.T) {
+	vm := MakeVM()
+	vm.TLAVar("fooString", "bar")
+	vm.TLACode("fooCode", "true")
+	_, err := vm.EvaluateAnonymousSnippet("test.jsonnet", `function (fooString, fooCode) { str: fooString, code: fooCode }`)
+	if err != nil {
+		t.Fatalf("unexpected error %v", err)
+	}
+	vm.TLAReset()
+	_, err = vm.EvaluateAnonymousSnippet("test.jsonnet", `function(fooString, fooCode) { str: fooString, code: fooCode }`)
+	if err == nil {
+		t.Fatalf("expected error, got nil")
+	}
+	if !strings.Contains(err.Error(), "Missing argument") {
+		t.Errorf("unexpected error %v", err)
+	}
+}
diff --git a/vm.go b/vm.go
index 4ba94d4..43c0f8c 100644
--- a/vm.go
+++ b/vm.go
@@ -95,6 +95,12 @@ func (vm *VM) ExtCode(key string, val string) {
 	vm.flushValueCache()
 }
 
+// ExtReset rests all external variables registered for this VM.
+func (vm *VM) ExtReset() {
+	vm.ext = make(vmExtMap)
+	vm.flushValueCache()
+}
+
 // TLAVar binds a Jsonnet top level argument to the given value.
 func (vm *VM) TLAVar(key string, val string) {
 	vm.tla[key] = vmExt{value: val, isCode: false}
@@ -109,6 +115,11 @@ func (vm *VM) TLACode(key string, val string) {
 	// Setting a TLA does not require flushing the cache - see above.
 }
 
+// TLAReset resets all TLAs registered for this VM.
+func (vm *VM) TLAReset() {
+	vm.tla = make(vmExtMap)
+}
+
 // Importer sets Importer to use during evaluation (import callback).
 func (vm *VM) Importer(i Importer) {
 	vm.importer = i