From 790f62b41543cfb65850359ab05a43829cf0ba18 Mon Sep 17 00:00:00 2001
From: "M. J. Fromberger" <michael.j.fromberger@gmail.com>
Date: Tue, 25 Jan 2022 18:05:46 -0800
Subject: [PATCH] handler: add decoding tests for custom JSON types

---
 handler/handler_test.go | 80 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/handler/handler_test.go b/handler/handler_test.go
index b61ecde6..6f057e8d 100644
--- a/handler/handler_test.go
+++ b/handler/handler_test.go
@@ -6,6 +6,8 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
+	"strconv"
 	"testing"
 
 	"github.com/creachadair/jrpc2"
@@ -71,6 +73,49 @@ func TestCheck(t *testing.T) {
 	}
 }
 
+// Verify that the wrappers constructed by FuncInfo.Wrap can properly decode
+// their arguments of different types and structure.
+func TestFuncInfo_wrapDecode(t *testing.T) {
+	tests := []struct {
+		fn   handler.Func
+		p    string
+		want interface{}
+	}{
+		// A positional handler should decode its argument from an array or an object.
+		{handler.NewPos(func(_ context.Context, z int) int { return z }, "arg"),
+			`[25]`, 25},
+		{handler.NewPos(func(_ context.Context, z int) int { return z }, "arg"),
+			`{"arg":109}`, 109},
+
+		// A type with custom marshaling should be properly handled.
+		{handler.NewPos(func(_ context.Context, b stringByte) byte { return byte(b) }, "arg"),
+			`["00111010"]`, byte(0x3a)},
+		{handler.NewPos(func(_ context.Context, b stringByte) byte { return byte(b) }, "arg"),
+			`{"arg":"10011100"}`, byte(0x9c)},
+		{handler.New(func(_ context.Context, v fauxStruct) int { return int(v) }),
+			`{"type":"thing","value":99}`, 99},
+
+		// Plain JSON should get its argument unmodified.
+		{handler.New(func(_ context.Context, v json.RawMessage) string { return string(v) }),
+			`{"x": true, "y": null}`, `{"x": true, "y": null}`},
+
+		// Npn-positional slice argument.
+		{handler.New(func(_ context.Context, ss []string) int { return len(ss) }),
+			`["a", "b", "c"]`, 3},
+	}
+	ctx := context.Background()
+	for _, test := range tests {
+		req := mustParseRequest(t,
+			fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"method":"x","params":%s}`, test.p))
+		got, err := test.fn(ctx, req)
+		if err != nil {
+			t.Errorf("Call %v failed: %v", test.fn, err)
+		} else if diff := cmp.Diff(test.want, got); diff != "" {
+			t.Errorf("Call %v: wrong result (-want, +got)\n%s", test.fn, diff)
+		}
+	}
+}
+
 // Verify that the Positional function correctly handles its cases.
 func TestPositional(t *testing.T) {
 	tests := []struct {
@@ -403,3 +448,38 @@ func mustParseRequest(t *testing.T, text string) *jrpc2.Request {
 	}
 	return req[0]
 }
+
+// stringByte is a byte with a custom JSON encoding. It expects a string of
+// decimal digits 1 and 0, e.g., "10011000" == 0x98.
+type stringByte byte
+
+func (s *stringByte) UnmarshalText(text []byte) error {
+	v, err := strconv.ParseUint(string(text), 2, 8)
+	if err != nil {
+		return err
+	}
+	*s = stringByte(v)
+	return nil
+}
+
+// fauxStruct is an integer with a custom JSON encoding. It expects an object:
+//
+//   {"type":"thing","value":<integer>}
+//
+type fauxStruct int
+
+func (s *fauxStruct) UnmarshalJSON(data []byte) error {
+	var tmp struct {
+		T string `json:"type"`
+		V *int   `json:"value"`
+	}
+	if err := json.Unmarshal(data, &tmp); err != nil {
+		return err
+	} else if tmp.T != "thing" {
+		return fmt.Errorf("unknown type %q", tmp.T)
+	} else if tmp.V == nil {
+		return errors.New("missing value")
+	}
+	*s = fauxStruct(*tmp.V)
+	return nil
+}