From 318642ac6f0846a018af6d812a013b08fdb70dac Mon Sep 17 00:00:00 2001
From: Nikita Pivkin <nikita.pivkin@smartforce.io>
Date: Fri, 8 Dec 2023 21:13:42 +0300
Subject: [PATCH] feat(terraform): add a method to replace the value in the
 context (#1504)

---
 pkg/terraform/context/context.go      | 10 ++++++++++
 pkg/terraform/context/context_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/pkg/terraform/context/context.go b/pkg/terraform/context/context.go
index 5b5d4d274..496aad1cb 100644
--- a/pkg/terraform/context/context.go
+++ b/pkg/terraform/context/context.go
@@ -77,6 +77,16 @@ func (c *Context) Set(val cty.Value, parts ...string) {
 	c.ctx.Variables[parts[0]] = v
 }
 
+func (c *Context) Replace(val cty.Value, path string) {
+	parts := strings.Split(path, ".")
+	if len(parts) == 0 {
+		return
+	}
+
+	delete(c.ctx.Variables, parts[0])
+	c.Set(val, parts...)
+}
+
 func mergeVars(src cty.Value, parts []string, value cty.Value) cty.Value {
 
 	if len(parts) == 0 {
diff --git a/pkg/terraform/context/context_test.go b/pkg/terraform/context/context_test.go
index 57da7498a..8185d7b98 100644
--- a/pkg/terraform/context/context_test.go
+++ b/pkg/terraform/context/context_test.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/hashicorp/hcl/v2"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/gocty"
 )
@@ -211,3 +212,27 @@ func Test_IsNotEmptyObject(t *testing.T) {
 		})
 	}
 }
+
+func TestReplace(t *testing.T) {
+	t.Run("replacement of an existing value", func(t *testing.T) {
+		underlying := &hcl.EvalContext{}
+		ctx := NewContext(underlying, nil)
+		ctx.SetByDot(cty.StringVal("some-value"), "my.value")
+		require.NotEqual(t, cty.NilVal, ctx.GetByDot("my.value"))
+		ctx.Replace(cty.NumberIntVal(-1), "my.value")
+		assert.Equal(t, cty.NumberIntVal(-1), ctx.GetByDot("my.value"))
+	})
+
+	t.Run("replacement of a non-existing value", func(t *testing.T) {
+		underlying := &hcl.EvalContext{}
+		ctx := NewContext(underlying, nil)
+		ctx.Replace(cty.NumberIntVal(-1), "my.value")
+		assert.Equal(t, cty.NumberIntVal(-1), ctx.GetByDot("my.value"))
+	})
+
+	t.Run("empty path", func(t *testing.T) {
+		underlying := &hcl.EvalContext{}
+		ctx := NewContext(underlying, nil)
+		ctx.Replace(cty.NumberIntVal(-1), "")
+	})
+}