From e9498a983f35578906bb732496a38efa5aa1e4e6 Mon Sep 17 00:00:00 2001
From: Alex Efros <powerman@powerman.name>
Date: Wed, 22 Apr 2020 15:55:01 +0300
Subject: [PATCH] feat: add WrapError

---
 jsonrpc2/errors.go       | 24 ++++++++++++++++++++++++
 jsonrpc2/errors_test.go  | 27 +++++++++++++++++++++++++++
 jsonrpc2/example_test.go | 20 +++++++-------------
 3 files changed, 58 insertions(+), 13 deletions(-)
 create mode 100644 jsonrpc2/errors_test.go

diff --git a/jsonrpc2/errors.go b/jsonrpc2/errors.go
index 4c930eb..f7b017e 100644
--- a/jsonrpc2/errors.go
+++ b/jsonrpc2/errors.go
@@ -3,6 +3,8 @@ package jsonrpc2
 import (
 	"encoding/json"
 	"fmt"
+	"io"
+	"net/rpc"
 	"strings"
 )
 
@@ -90,3 +92,25 @@ func (e *Error) Error() string {
 	}
 	return string(buf)
 }
+
+type wrapError struct {
+	err *Error
+}
+
+// WrapError handles any error returned by Client.Call() by wrapping Error
+// returned by ServerError or returning non-ServerError errors as is.
+// Wrapped Error is stringified to "<code> <message>" instead of JSON.
+func WrapError(err error) error {
+	if err == nil || err == rpc.ErrShutdown || err == io.ErrUnexpectedEOF {
+		return err
+	}
+	return wrapError{err: ServerError(err)}
+}
+
+func (e wrapError) Error() string {
+	return fmt.Sprintf("%d %s", e.err.Code, e.err.Message)
+}
+
+func (e wrapError) Unwrap() error {
+	return e.err
+}
diff --git a/jsonrpc2/errors_test.go b/jsonrpc2/errors_test.go
new file mode 100644
index 0000000..106c66f
--- /dev/null
+++ b/jsonrpc2/errors_test.go
@@ -0,0 +1,27 @@
+package jsonrpc2
+
+import (
+	"io"
+	"net/rpc"
+	"testing"
+)
+
+func TestWrapError(t *testing.T) {
+	tests := []struct {
+		err error
+	}{
+		{nil},
+		{rpc.ErrShutdown},
+		{io.ErrUnexpectedEOF},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run("", func(t *testing.T) {
+			err := WrapError(tc.err)
+			if err != tc.err {
+				t.Errorf("got %v, want %v", err, tc.err)
+			}
+		})
+	}
+
+}
diff --git a/jsonrpc2/example_test.go b/jsonrpc2/example_test.go
index ae2a205..4731398 100644
--- a/jsonrpc2/example_test.go
+++ b/jsonrpc2/example_test.go
@@ -165,20 +165,14 @@ func Example() {
 	clientHTTP.Call("ExampleSvc.FullName3", NameArg{"First", "Last"}, nil)
 
 	// Correct error handling.
-	err = clientTCP.Call("ExampleSvc.Err1", nil, nil)
-	if err == rpc.ErrShutdown || err == io.ErrUnexpectedEOF {
-		fmt.Printf("Err1(): %q\n", err)
-	} else if err != nil {
-		rpcerr := jsonrpc2.ServerError(err)
-		fmt.Printf("Err1(): code=%d msg=%q data=%v\n", rpcerr.Code, rpcerr.Message, rpcerr.Data)
-	}
+	err = jsonrpc2.WrapError(clientTCP.Call("ExampleSvc.Err1", nil, nil))
+	fmt.Printf("Err1(): %q\n", err)
 
-	err = clientCustomHTTP.Call("ExampleSvc.Err2", nil, nil)
-	if err == rpc.ErrShutdown || err == io.ErrUnexpectedEOF {
-		fmt.Printf("Err2(): %q\n", err)
-	} else if err != nil {
-		rpcerr := jsonrpc2.ServerError(err)
+	err = jsonrpc2.WrapError(clientCustomHTTP.Call("ExampleSvc.Err2", nil, nil))
+	if rpcerr := new(jsonrpc2.Error); errors.As(err, &rpcerr) {
 		fmt.Printf("Err2(): code=%d msg=%q data=%v\n", rpcerr.Code, rpcerr.Message, rpcerr.Data)
+	} else if err != nil {
+		fmt.Printf("Err2(): %q\n", err)
 	}
 
 	err = clientHTTP.Call("ExampleSvc.Err3", nil, nil)
@@ -195,7 +189,7 @@ func Example() {
 	// MapLen({a:10,b:20,c:30})=3
 	// FullName2(): Remote IP is 127.0.0.1
 	// FullName3(): Remote IP is 127.0.0.1
-	// Err1(): code=-32000 msg="some issue" data=<nil>
+	// Err1(): "-32000 some issue"
 	// Err2(): code=-32603 msg="bad HTTP Status: 415 Unsupported Media Type" data=<nil>
 	// Err3(): code=42 msg="some issue" data=[one two]
 }