From 8b7188a46af2b5ecf9879f4ccd45f543393d235a Mon Sep 17 00:00:00 2001
From: Manik Rana <Manikrana54@gmail.com>
Date: Thu, 18 Jan 2024 02:54:55 +0530
Subject: [PATCH] tests: add tests for `go/json2` (#14964)

Signed-off-by: Manik Rana <manikrana54@gmail.com>
---
 go/json2/marshal_test.go   | 23 +++++++++++++----
 go/json2/unmarshal_test.go | 51 +++++++++++++++++++++++++++++++++++---
 2 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/go/json2/marshal_test.go b/go/json2/marshal_test.go
index 96b7f508d73..b155126fb17 100644
--- a/go/json2/marshal_test.go
+++ b/go/json2/marshal_test.go
@@ -19,6 +19,9 @@ package json2
 import (
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
 )
@@ -29,11 +32,21 @@ func TestMarshalPB(t *testing.T) {
 		Type: querypb.Type_VARCHAR,
 	}
 	b, err := MarshalPB(col)
-	if err != nil {
-		t.Fatal(err)
-	}
+
+	require.NoErrorf(t, err, "MarshalPB(%+v) error", col)
 	want := "{\"name\":\"c1\",\"type\":\"VARCHAR\"}"
-	if string(b) != want {
-		t.Errorf("MarshalPB(col): %q, want %q", b, want)
+	assert.Equalf(t, want, string(b), "MarshalPB(%+v)", col)
+}
+
+func TestMarshalIndentPB(t *testing.T) {
+	col := &vschemapb.Column{
+		Name: "c1",
+		Type: querypb.Type_VARCHAR,
 	}
+	indent := "  "
+	b, err := MarshalIndentPB(col, indent)
+
+	require.NoErrorf(t, err, "MarshalIndentPB(%+v, %q) error", col, indent)
+	want := "{\n  \"name\": \"c1\",\n  \"type\": \"VARCHAR\"\n}"
+	assert.Equal(t, want, string(b), "MarshalIndentPB(%+v, %q)", col, indent)
 }
diff --git a/go/json2/unmarshal_test.go b/go/json2/unmarshal_test.go
index 9b6a6af1ca2..ff18a29def8 100644
--- a/go/json2/unmarshal_test.go
+++ b/go/json2/unmarshal_test.go
@@ -17,7 +17,14 @@ limitations under the License.
 package json2
 
 import (
+	"fmt"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/types/known/emptypb"
 )
 
 func TestUnmarshal(t *testing.T) {
@@ -37,14 +44,50 @@ func TestUnmarshal(t *testing.T) {
 		err: "",
 	}}
 	for _, tcase := range tcases {
-		out := make(map[string]any)
+		out := make(map[string]interface{})
 		err := Unmarshal([]byte(tcase.in), &out)
+
 		got := ""
 		if err != nil {
 			got = err.Error()
 		}
-		if got != tcase.err {
-			t.Errorf("Unmarshal(%v) err: %v, want %v", tcase.in, got, tcase.err)
-		}
+		assert.Equal(t, tcase.err, got, "Unmarshal(%v) err", tcase.in)
+	}
+}
+
+func TestUnmarshalProto(t *testing.T) {
+	protoData := &emptypb.Empty{}
+	protoJSONData, err := protojson.Marshal(protoData)
+	assert.Nil(t, err, "protojson.Marshal error")
+
+	tcase := struct {
+		in  string
+		out *emptypb.Empty
+	}{
+		in:  string(protoJSONData),
+		out: &emptypb.Empty{},
+	}
+
+	err = Unmarshal([]byte(tcase.in), tcase.out)
+
+	assert.Nil(t, err, "Unmarshal(%v) protobuf message", tcase.in)
+	assert.Equal(t, protoData, tcase.out, "Unmarshal(%v) protobuf message result", tcase.in)
+}
+
+func TestAnnotate(t *testing.T) {
+	tcases := []struct {
+		data []byte
+		err  error
+	}{
+		{
+			data: []byte("invalid JSON"),
+			err:  fmt.Errorf("line: 1, position 1: invalid character 'i' looking for beginning of value"),
+		},
+	}
+
+	for _, tcase := range tcases {
+		err := annotate(tcase.data, tcase.err)
+
+		require.Equal(t, tcase.err, err, "annotate(%s, %v) error", string(tcase.data), tcase.err)
 	}
 }