From 82e8930f9ab0a6b5348258ca21b68d72b2fa7591 Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@gmail.com>
Date: Sun, 29 May 2022 16:53:20 -0400
Subject: [PATCH] assert: tests for golden variables

---
 assert/assert_ext_test.go | 72 ++++++++++++++++++++++++++++++++-------
 internal/source/source.go |  2 +-
 internal/source/update.go | 27 ++++++++++-----
 3 files changed, 80 insertions(+), 21 deletions(-)

diff --git a/assert/assert_ext_test.go b/assert/assert_ext_test.go
index 88b0421..5903f70 100644
--- a/assert/assert_ext_test.go
+++ b/assert/assert_ext_test.go
@@ -1,7 +1,11 @@
 package assert_test
 
 import (
-	"fmt"
+	"go/parser"
+	"go/token"
+	"io/ioutil"
+	"runtime"
+	"strings"
 	"testing"
 
 	"gotest.tools/v3/assert"
@@ -9,32 +13,56 @@ import (
 )
 
 func TestEqual_WithGoldenUpdate(t *testing.T) {
-	t.Run("assert failed with update=false", func(t *testing.T) {
+	t.Run("assert failed with -update=false", func(t *testing.T) {
 		ft := &fakeTestingT{}
 		actual := `not this value`
 		assert.Equal(ft, actual, expectedOne)
 		assert.Assert(t, ft.failNowed)
 	})
 
-	t.Run("value is updated when -update=true", func(t *testing.T) {
+	t.Run("var is updated when -update=true", func(t *testing.T) {
 		patchUpdate(t)
-		ft := &fakeTestingT{}
+		t.Cleanup(func() {
+			resetVariable(t, "expectedOne", "")
+		})
 
 		actual := `this is the
 actual value
-that we are testing against`
-		assert.Equal(ft, actual, expectedOne)
+that we are testing
+`
+		assert.Equal(t, actual, expectedOne)
 
-		// reset
-		fmt.Println("WHHHHHHHHHHY")
-		assert.Equal(ft, "\n\n\n", expectedOne)
-	})
-}
+		raw, err := ioutil.ReadFile(fileName(t))
+		assert.NilError(t, err)
 
-var expectedOne = `
+		expected := "var expectedOne = `this is the\nactual value\nthat we are testing\n`"
+		assert.Assert(t, strings.Contains(string(raw), expected), "actual=%v", string(raw))
+	})
 
+	t.Run("const is updated when -update=true", func(t *testing.T) {
+		patchUpdate(t)
+		t.Cleanup(func() {
+			resetVariable(t, "expectedTwo", "")
+		})
 
+		actual := `this is the new
+expected value
 `
+		assert.Equal(t, actual, expectedTwo)
+
+		raw, err := ioutil.ReadFile(fileName(t))
+		assert.NilError(t, err)
+
+		expected := "const expectedTwo = `this is the new\nexpected value\n`"
+		assert.Assert(t, strings.Contains(string(raw), expected), "actual=%v", string(raw))
+	})
+}
+
+// expectedOne is updated by running the tests with -update
+var expectedOne = ``
+
+// expectedTwo is updated by running the tests with -update
+const expectedTwo = ``
 
 func patchUpdate(t *testing.T) {
 	source.Update = true
@@ -43,6 +71,26 @@ func patchUpdate(t *testing.T) {
 	})
 }
 
+func fileName(t *testing.T) string {
+	t.Helper()
+	_, filename, _, ok := runtime.Caller(1)
+	assert.Assert(t, ok, "failed to get call stack")
+	return filename
+}
+
+func resetVariable(t *testing.T, varName string, value string) {
+	t.Helper()
+	_, filename, _, ok := runtime.Caller(1)
+	assert.Assert(t, ok, "failed to get call stack")
+
+	fileset := token.NewFileSet()
+	astFile, err := parser.ParseFile(fileset, filename, nil, parser.AllErrors|parser.ParseComments)
+	assert.NilError(t, err)
+
+	err = source.UpdateVariable(filename, fileset, astFile, varName, value)
+	assert.NilError(t, err, "failed to reset file")
+}
+
 type fakeTestingT struct {
 	failNowed bool
 	failed    bool
diff --git a/internal/source/source.go b/internal/source/source.go
index 453bee4..a3f7008 100644
--- a/internal/source/source.go
+++ b/internal/source/source.go
@@ -75,7 +75,7 @@ func scanToLine(fileset *token.FileSet, node ast.Node, lineNum int) ast.Node {
 	return matchedNode
 }
 
-func getCallExprArgs(fileset *token.FileSet, astFile *ast.File, line int) ([]ast.Expr, error) {
+func getCallExprArgs(fileset *token.FileSet, astFile ast.Node, line int) ([]ast.Expr, error) {
 	node, err := getNodeAtLine(fileset, astFile, line)
 	switch {
 	case err != nil:
diff --git a/internal/source/update.go b/internal/source/update.go
index 1b669f6..bd9678b 100644
--- a/internal/source/update.go
+++ b/internal/source/update.go
@@ -43,8 +43,6 @@ func UpdateExpectedValue(stackIndex int, x, y interface{}) error {
 		return fmt.Errorf("failed to parse source file %s: %w", filename, err)
 	}
 
-	debug("before modification: %v", debugFormatNode{astFile})
-
 	expr, err := getCallExprArgs(fileset, astFile, line)
 	if err != nil {
 		return fmt.Errorf("call from %s:%d: %w", filename, line, err)
@@ -68,6 +66,23 @@ func UpdateExpectedValue(stackIndex int, x, y interface{}) error {
 		value = y
 	}
 
+	strValue, ok := value.(string)
+	if !ok {
+		debug("value must be type string, got %T", value)
+		return ErrNotFound
+	}
+	return UpdateVariable(filename, fileset, astFile, varName, strValue)
+}
+
+// UpdateVariable writes to filename the contents of astFile with the value of
+// the variable updated to value.
+func UpdateVariable(
+	filename string,
+	fileset *token.FileSet,
+	astFile *ast.File,
+	varName string,
+	value string,
+) error {
 	obj := astFile.Scope.Objects[varName]
 	if obj == nil {
 		return ErrNotFound
@@ -87,15 +102,11 @@ func UpdateExpectedValue(stackIndex int, x, y interface{}) error {
 		return ErrNotFound
 	}
 
-	// TODO: allow a function to wrap the string literal
 	spec.Values[0] = &ast.BasicLit{
-		Kind: token.STRING,
-		// TODO: safer
-		Value: "`" + value.(string) + "`",
+		Kind:  token.STRING,
+		Value: "`" + value + "`",
 	}
 
-	debug("after modification: %v", debugFormatNode{astFile})
-
 	var buf bytes.Buffer
 	if err := format.Node(&buf, fileset, astFile); err != nil {
 		return fmt.Errorf("failed to format file after update: %w", err)