From 2941ee12eb653ae6965af2c79213b60a2e875008 Mon Sep 17 00:00:00 2001
From: dfawley <dfawley@google.com>
Date: Mon, 11 Dec 2017 09:02:19 -0800
Subject: [PATCH] codes: Add UnmarshalJSON support to Code type (#1720)

---
 codes/codes.go      | 41 +++++++++++++++++++++++++++++
 codes/codes_test.go | 64 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)
 create mode 100644 codes/codes_test.go

diff --git a/codes/codes.go b/codes/codes.go
index 01382921..f3719d56 100644
--- a/codes/codes.go
+++ b/codes/codes.go
@@ -19,6 +19,9 @@
 // Package codes defines the canonical error codes used by gRPC. It is
 // consistent across various languages.
 package codes // import "google.golang.org/grpc/codes"
+import (
+	"fmt"
+)
 
 // A Code is an unsigned 32-bit error code as defined in the gRPC spec.
 type Code uint32
@@ -140,3 +143,41 @@ const (
 	// DataLoss indicates unrecoverable data loss or corruption.
 	DataLoss Code = 15
 )
+
+var strToCode = map[string]Code{
+	`"OK"`: OK,
+	`"CANCELLED"`:/* [sic] */ Canceled,
+	`"UNKNOWN"`:             Unknown,
+	`"INVALID_ARGUMENT"`:    InvalidArgument,
+	`"DEADLINE_EXCEEDED"`:   DeadlineExceeded,
+	`"NOT_FOUND"`:           NotFound,
+	`"ALREADY_EXISTS"`:      AlreadyExists,
+	`"PERMISSION_DENIED"`:   PermissionDenied,
+	`"RESOURCE_EXHAUSTED"`:  ResourceExhausted,
+	`"FAILED_PRECONDITION"`: FailedPrecondition,
+	`"ABORTED"`:             Aborted,
+	`"OUT_OF_RANGE"`:        OutOfRange,
+	`"UNIMPLEMENTED"`:       Unimplemented,
+	`"INTERNAL"`:            Internal,
+	`"UNAVAILABLE"`:         Unavailable,
+	`"DATA_LOSS"`:           DataLoss,
+	`"UNAUTHENTICATED"`:     Unauthenticated,
+}
+
+// UnmarshalJSON unmarshals b into the Code.
+func (c *Code) UnmarshalJSON(b []byte) error {
+	// From json.Unmarshaler: By convention, to approximate the behavior of
+	// Unmarshal itself, Unmarshalers implement UnmarshalJSON([]byte("null")) as
+	// a no-op.
+	if string(b) == "null" {
+		return nil
+	}
+	if c == nil {
+		return fmt.Errorf("nil receiver passed to UnmarshalJSON")
+	}
+	if jc, ok := strToCode[string(b)]; ok {
+		*c = jc
+		return nil
+	}
+	return fmt.Errorf("invalid code: %q", string(b))
+}
diff --git a/codes/codes_test.go b/codes/codes_test.go
new file mode 100644
index 00000000..1e3b9918
--- /dev/null
+++ b/codes/codes_test.go
@@ -0,0 +1,64 @@
+/*
+ *
+ * Copyright 2017 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package codes
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	cpb "google.golang.org/genproto/googleapis/rpc/code"
+)
+
+func TestUnmarshalJSON(t *testing.T) {
+	for s, v := range cpb.Code_value {
+		want := Code(v)
+		var got Code
+		if err := got.UnmarshalJSON([]byte(`"` + s + `"`)); err != nil || got != want {
+			t.Errorf("got.UnmarshalJSON(%q) = %v; want <nil>.  got=%v; want %v", s, err, got, want)
+		}
+	}
+}
+
+func TestJSONUnmarshal(t *testing.T) {
+	var got []Code
+	want := []Code{OK, NotFound, Internal, Canceled}
+	in := `["OK", "NOT_FOUND", "INTERNAL", "CANCELLED"]`
+	err := json.Unmarshal([]byte(in), &got)
+	if err != nil || !reflect.DeepEqual(got, want) {
+		t.Fatalf("json.Unmarshal(%q, &got) = %v; want <nil>.  got=%v; want %v", in, err, got, want)
+	}
+}
+
+func TestUnmarshalJSON_NilReceiver(t *testing.T) {
+	var got *Code
+	in := OK.String()
+	if err := got.UnmarshalJSON([]byte(in)); err == nil {
+		t.Errorf("got.UnmarshalJSON(%q) = nil; want <non-nil>.  got=%v", in, got)
+	}
+}
+
+func TestUnmarshalJSON_UnknownInput(t *testing.T) {
+	var got Code
+	for _, in := range [][]byte{[]byte(""), []byte("xxx"), []byte("Code(17)"), nil} {
+		if err := got.UnmarshalJSON([]byte(in)); err == nil {
+			t.Errorf("got.UnmarshalJSON(%q) = nil; want <non-nil>.  got=%v", in, got)
+		}
+	}
+}