Return grpc server error code

bavix · Sep 2, 2023 · a96033a · a96033a
1 parent f645650
commit a96033a
Show file tree

Hide file tree

Showing 8 changed files with 169 additions and 14 deletions.
diff --git a/Readme.md b/Readme.md
@@ -76,7 +76,8 @@ Stub Format is JSON text format. It has a skeleton as follows:
     "data":{
       // put result fields here
     },
-    "error":"<error message>" // Optional. if you want to return error instead.
+    "error":"<error message>", // Optional. if you want to return error instead.
+    "code":"<response code>" // Optional. Grpc response code. if code !=0  return error instead.
   }
 }
 ```

diff --git a/example/simple/client/main.go b/example/simple/client/main.go
@@ -6,9 +6,12 @@ import (
 	"os"
 	"time"
 
-	pb "github.com/bavix/gripmock/protogen/example/simple"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
+
+	pb "github.com/bavix/gripmock/protogen/example/simple"
 )
 
 //nolint:gomnd
@@ -68,4 +71,28 @@ func main() {
 		log.Fatalf("grpc server returned code: %d, expected code: %d", r.ReturnCode, 3)
 	}
 	log.Printf("Greeting: %s (return code %d)", r.Message, r.ReturnCode)
+
+	name = "error"
+	r, err = c.SayHello(context.Background(), &pb.Request{Name: name})
+	if err == nil {
+		log.Fatalf("Expected error, but return %d", r.ReturnCode)
+	}
+	log.Printf("Greeting error: %s", err)
+
+	name = "error_code"
+	r, err = c.SayHello(context.Background(), &pb.Request{Name: name})
+	if err == nil {
+		log.Fatalf("Expected error, but return %d", r.ReturnCode)
+	}
+
+	s, ok := status.FromError(err)
+	if !ok {
+		log.Fatalf("Expected to get error status: %v", err)
+	}
+
+	if s.Code() != codes.InvalidArgument {
+		log.Fatalf("Expected to get error status %d, got: %d", codes.InvalidArgument, s.Code())
+	}
+
+	log.Printf("Greeting error: %s, code: %d", err, s.Code())
 }
diff --git a/example/simple/stub/errors.yml b/example/simple/stub/errors.yml
@@ -0,0 +1,16 @@
+---
+- service: Gripmock
+  method: SayHello
+  input:
+    equals:
+      name: error
+  output:
+    error: test_error
+- service: Gripmock
+  method: SayHello
+  input:
+    equals:
+      name: error_code
+  output:
+    error: test_error_code
+    code: 3
diff --git a/pkg/sdk/client.go b/pkg/sdk/client.go
@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+
+	"google.golang.org/grpc/codes"
 )
 
 type StubApiClient struct {
@@ -27,9 +29,10 @@ type Payload struct {
 type Response struct {
 	Data  interface{} `json:"data"`
 	Error string      `json:"error"`
+	Code  *codes.Code `json:"code,omitempty"`
 }
 
-func (c *StubApiClient) Search(ctx context.Context, payload Payload) (any, error) {
+func (c *StubApiClient) Search(ctx context.Context, payload Payload) (*Response, error) {
 	postBody, err := json.Marshal(payload)
 	if err != nil {
 		return nil, err
@@ -63,11 +66,5 @@ func (c *StubApiClient) Search(ctx context.Context, payload Payload) (any, error
 		return nil, err
 	}
 
-	if result.Error != "" {
-		//fixme
-		//nolint:goerr113
-		return nil, fmt.Errorf(result.Error)
-	}
-
-	return result.Data, nil
+	return result, nil
 }
diff --git a/pkg/storage/stubs.go b/pkg/storage/stubs.go
@@ -5,6 +5,7 @@ import (
 	"sync"
 
 	"github.com/google/uuid"
+	"google.golang.org/grpc/codes"
 )
 
 var ErrServiceNotFound = errors.New("service not found")
@@ -27,6 +28,7 @@ type Input struct {
 type Output struct {
 	Data  map[string]interface{} `json:"data"`
 	Error string                 `json:"error"`
+	Code  *codes.Code            `json:"code,omitempty"`
 }
 
 type storage struct {

diff --git a/protoc-gen-gripmock/server.tmpl b/protoc-gen-gripmock/server.tmpl
@@ -17,6 +17,8 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
 	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/bavix/gripmock/pkg/sdk"
 )
@@ -153,7 +155,7 @@ type response struct {
 
 func findStub(ctx context.Context, service, method string, in, out protoreflect.ProtoMessage) error {
     api := sdk.NewStubApiClient(fmt.Sprintf("http://localhost%s", HTTP_PORT), http.DefaultClient)
-	resp, err := api.Search(ctx, sdk.Payload{
+	respRPC, err := api.Search(ctx, sdk.Payload{
         Service: service,
         Method:  method,
         Data:    in,
@@ -162,11 +164,22 @@ func findStub(ctx context.Context, service, method string, in, out protoreflect.
         return err
     }
 
-	data, err := json.Marshal(resp)
+    if respRPC.Error != "" || respRPC.Code != nil {
+		if respRPC.Code == nil {
+			abortedCode := codes.Aborted
+			respRPC.Code = &abortedCode
+		}
+
+		if *respRPC.Code != codes.OK {
+			return status.Error(*respRPC.Code, respRPC.Error)
+		}
+	}
+
+	data, err := json.Marshal(respRPC.Data)
 	if err != nil {
 	    return err
 	}
 
 	return jsonpb.Unmarshal(data, out)
 }
-{{ end }}
+{{ end }}
diff --git a/stub/api_test.go b/stub/api_test.go
@@ -230,6 +230,105 @@ func TestStub(t *testing.T) {
 			handler: api.AddHandle,
 			expect:  `["3f68f410-bb58-49ad-b679-23f2ed690c1d","6da11d72-c0db-4075-9e72-31d61ffd0483"]`,
 		},
+
+		{
+			name: "add error stub with result code contains",
+			mock: func() *http.Request {
+				payload := `{
+								"id": "cda7321b-9241-4a58-9cbf-0603e0146542",
+								"service": "ErrorStabWithCode",
+								"method":"TestMethod",
+								"input":{
+									"contains":{
+												"key": "value",
+												"greetings": {
+													"hola": "mundo",
+													"merhaba": "dunya"
+												},
+												"cities": ["Istanbul", "Jakarta"]
+									}
+								},
+								"output":{
+									"error":"error msg",
+                                    "code": 3
+								}
+							}`
+				return httptest.NewRequest(http.MethodPost, "/api/stubs", bytes.NewReader([]byte(payload)))
+			},
+			handler: api.AddHandle,
+			expect:  `["cda7321b-9241-4a58-9cbf-0603e0146542"]`,
+		},
+		{
+			name: "find error stub with result code contains",
+			mock: func() *http.Request {
+				payload := `{
+						"service": "ErrorStabWithCode",
+						"method":"TestMethod",
+						"data":{
+								"key": "value",
+								"anotherKey": "anotherValue",
+								"greetings": {
+									"hola": "mundo",
+									"merhaba": "dunya",
+									"hello": "world"
+								},
+								"cities": ["Istanbul", "Jakarta", "Winterfell"]
+						}
+					}`
+				return httptest.NewRequest(http.MethodPost, "/api/stubs/search", bytes.NewReader([]byte(payload)))
+			},
+			handler: api.SearchHandle,
+			expect:  "{\"data\":null,\"error\":\"error msg\",\"code\":3}\n",
+		},
+
+		{
+			name: "add error stub without result code contains",
+			mock: func() *http.Request {
+				payload := `{
+								"id": "6d5ec9a6-94a7-4f23-b5ea-b04a37796adb",
+								"service": "ErrorStab",
+								"method":"TestMethod",
+								"input":{
+									"contains":{
+												"key": "value",
+												"greetings": {
+													"hola": "mundo",
+													"merhaba": "dunya"
+												},
+												"cities": ["Istanbul", "Jakarta"]
+									}
+								},
+								"output":{
+									"error":"error msg"
+								}
+							}`
+				return httptest.NewRequest(http.MethodPost, "/api/stubs", bytes.NewReader([]byte(payload)))
+			},
+			handler: api.AddHandle,
+			expect:  `["6d5ec9a6-94a7-4f23-b5ea-b04a37796adb"]`,
+		},
+		{
+			name: "find error stub without result code contains",
+			mock: func() *http.Request {
+				payload := `{
+						"service": "ErrorStab",
+						"method":"TestMethod",
+						"data":{
+								"key": "value",
+								"anotherKey": "anotherValue",
+								"greetings": {
+									"hola": "mundo",
+									"merhaba": "dunya",
+									"hello": "world"
+								},
+								"cities": ["Istanbul", "Jakarta", "Winterfell"]
+						}
+					}`
+				return httptest.NewRequest(http.MethodPost, "/api/stubs/search", bytes.NewReader([]byte(payload)))
+			},
+			handler: api.SearchHandle,
+			expect:  "{\"data\":null,\"error\":\"error msg\"}\n",
+		},
 		{
 			name: "find nested stub contains",
 			mock: func() *http.Request {

diff --git a/stub/stub.go b/stub/stub.go
@@ -83,7 +83,7 @@ func validateStub(stub *storage.Stub) error {
 
 	// TODO: validate all input case
 
-	if stub.Output.Error == "" && stub.Output.Data == nil {
+	if stub.Output.Error == "" && stub.Output.Data == nil && stub.Output.Code == nil {
 		//fixme
 		//nolint:goerr113
 		return fmt.Errorf("output can't be empty")