From b084dcea95eeb956ca2314c6123e088e429b8953 Mon Sep 17 00:00:00 2001
From: "Zhongduo Lin (Jimmy)" <zhongduo@google.com>
Date: Mon, 28 Sep 2020 16:49:03 -0400
Subject: [PATCH] feat(pubsub/pstest)!: Add status code to error injection
 (#2926)

A status code is added into `errorInjectionReactor`, this is necessary
as the retry setting for some calls depends on the error code. Also
the `errorInjectionReactor` is not exported. Note that the previous
version is same as having an `Unknown` code as the server will always
wrap the error to a proper status error.

BREAKING CHANGE: this will invalidate the previous exported
ErrorInjectionReactor struct and WithErrorInjection function.
---
 pubsub/pstest/fake.go      | 20 +++++++++++---------
 pubsub/pstest/fake_test.go | 13 +++++++++++--
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/pubsub/pstest/fake.go b/pubsub/pstest/fake.go
index 76620df0b30e..9b330874fbdd 100644
--- a/pubsub/pstest/fake.go
+++ b/pubsub/pstest/fake.go
@@ -1171,20 +1171,22 @@ func (s *GServer) runReactor(req interface{}, funcName string, defaultObj interf
 	return false, nil, nil
 }
 
-// ErrorInjectionReactor is a reactor to inject an error message
-type ErrorInjectionReactor struct {
-	errMsg string
+// errorInjectionReactor is a reactor to inject an error message with status code.
+type errorInjectionReactor struct {
+	code codes.Code
+	msg  string
 }
 
-// React simply returns an error with defined error message.
-func (e *ErrorInjectionReactor) React(_ interface{}) (handled bool, ret interface{}, err error) {
-	return true, nil, fmt.Errorf(e.errMsg)
+// React simply returns an error with defined error message and status code.
+func (e *errorInjectionReactor) React(_ interface{}) (handled bool, ret interface{}, err error) {
+	return true, nil, status.Errorf(e.code, e.msg)
 }
 
-// WithErrorInjection creates a ServerReactorOption that injects error for a certain function.
-func WithErrorInjection(funcName string, errMsg string) ServerReactorOption {
+// WithErrorInjection creates a ServerReactorOption that injects error with defined status code and
+// message for a certain function.
+func WithErrorInjection(funcName string, code codes.Code, msg string) ServerReactorOption {
 	return ServerReactorOption{
 		FuncName: funcName,
-		Reactor:  &ErrorInjectionReactor{errMsg: errMsg},
+		Reactor:  &errorInjectionReactor{code: code, msg: msg},
 	}
 }
diff --git a/pubsub/pstest/fake_test.go b/pubsub/pstest/fake_test.go
index 454a54a763fc..1ca043342587 100644
--- a/pubsub/pstest/fake_test.go
+++ b/pubsub/pstest/fake_test.go
@@ -951,15 +951,19 @@ func TestErrorInjection(t *testing.T) {
 	testcases := []struct {
 		funcName string
 		param    interface{}
+		code     codes.Code
 	}{
 		{
 			funcName: "CreateTopic",
+			code:     codes.Internal,
 		},
 		{
 			funcName: "GetTopic",
+			code:     codes.Aborted,
 		},
 		{
 			funcName: "UpdateTopic",
+			code:     codes.DeadlineExceeded,
 		},
 		{
 			funcName: "ListTopics",
@@ -1010,8 +1014,13 @@ func TestErrorInjection(t *testing.T) {
 	for _, tc := range testcases {
 		ctx := context.TODO()
 		errMsg := "error-injection-" + tc.funcName
+		// set error code to unknown unless specified
+		ec := codes.Unknown
+		if tc.code != codes.OK {
+			ec = tc.code
+		}
 		opts := []ServerReactorOption{
-			WithErrorInjection(tc.funcName, errMsg),
+			WithErrorInjection(tc.funcName, ec, errMsg),
 		}
 		_, _, server, cleanup := newFake(ctx, t, opts...)
 		defer cleanup()
@@ -1033,7 +1042,7 @@ func TestErrorInjection(t *testing.T) {
 		ret := reflect.ValueOf(&server.GServer).MethodByName(tc.funcName).Call([]reflect.Value{reflect.ValueOf(ctx), req})
 
 		got := ret[1].Interface().(error)
-		if got == nil || !strings.Contains(got.Error(), errMsg) {
+		if got == nil || status.Code(got) != ec || !strings.Contains(got.Error(), errMsg) {
 			t.Errorf("Got error does not contain the right key %v", got)
 		}
 	}