Replace usage of Code()

Signed-off-by: Daniel Nephin <[email protected]>

balancer/roundrobin/roundrobin_test.go

@@ -30,6 +30,7 @@ import (
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/codes"
 	_ "google.golang.org/grpc/grpclog/glogger"
+	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/resolver/manual"
@@ -111,7 +112,7 @@ func TestOneBackend(t *testing.T) {
 	// The first RPC should fail because there's no address.
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 	defer cancel()
-	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || grpc.Code(err) != codes.DeadlineExceeded {
+	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -143,7 +144,7 @@ func TestBackendsRoundRobin(t *testing.T) {
 	// The first RPC should fail because there's no address.
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 	defer cancel()
-	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || grpc.Code(err) != codes.DeadlineExceeded {
+	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -202,7 +203,7 @@ func TestAddressesRemoved(t *testing.T) {
 	// The first RPC should fail because there's no address.
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 	defer cancel()
-	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || grpc.Code(err) != codes.DeadlineExceeded {
+	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -216,7 +217,7 @@ func TestAddressesRemoved(t *testing.T) {
 	for i := 0; i < 1000; i++ {
 		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
 		defer cancel()
-		if _, err := testc.EmptyCall(ctx, &testpb.Empty{}, grpc.FailFast(false)); grpc.Code(err) == codes.DeadlineExceeded {
+		if _, err := testc.EmptyCall(ctx, &testpb.Empty{}, grpc.FailFast(false)); internal.Code(err) == codes.DeadlineExceeded {
 			return
 		}
 		time.Sleep(time.Millisecond)
@@ -248,7 +249,7 @@ func TestCloseWithPendingRPC(t *testing.T) {
 			defer wg.Done()
 			// This RPC blocks until cc is closed.
 			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-			if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); grpc.Code(err) == codes.DeadlineExceeded {
+			if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); internal.Code(err) == codes.DeadlineExceeded {
 				t.Errorf("RPC failed because of deadline after cc is closed; want error the client connection is closing")
 			}
 			cancel()
@@ -278,7 +279,7 @@ func TestNewAddressWhileBlocking(t *testing.T) {
 	// The first RPC should fail because there's no address.
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 	defer cancel()
-	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || grpc.Code(err) != codes.DeadlineExceeded {
+	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -327,7 +328,7 @@ func TestOneServerDown(t *testing.T) {
 	// The first RPC should fail because there's no address.
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 	defer cancel()
-	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || grpc.Code(err) != codes.DeadlineExceeded {
+	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -425,7 +426,7 @@ func TestAllServersDown(t *testing.T) {
 	// The first RPC should fail because there's no address.
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 	defer cancel()
-	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || grpc.Code(err) != codes.DeadlineExceeded {
+	if _, err := testc.EmptyCall(ctx, &testpb.Empty{}); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -469,7 +470,7 @@ func TestAllServersDown(t *testing.T) {
 	}
 	time.Sleep(100 * time.Millisecond)
 	for i := 0; i < 1000; i++ {
-		if _, err := testc.EmptyCall(context.Background(), &testpb.Empty{}); grpc.Code(err) == codes.Unavailable {
+		if _, err := testc.EmptyCall(context.Background(), &testpb.Empty{}); internal.Code(err) == codes.Unavailable {
 			return
 		}
 		time.Sleep(time.Millisecond)

balancer_test.go

@@ -34,6 +34,7 @@ import (
 
 	// V1 balancer tests use passthrough resolver instead of dns.
 	// TODO(bar) remove this when removing v1 balaner entirely.
+	"google.golang.org/grpc/internal"
 	_ "google.golang.org/grpc/resolver/passthrough"
 )
 
@@ -252,7 +253,7 @@ func TestCloseWithPendingRPC(t *testing.T) {
 	// Loop until the above update applies.
 	for {
 		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
-		if err := Invoke(ctx, "/foo/bar", &expectedRequest, &reply, cc, FailFast(false)); Code(err) == codes.DeadlineExceeded {
+		if err := Invoke(ctx, "/foo/bar", &expectedRequest, &reply, cc, FailFast(false)); internal.Code(err) == codes.DeadlineExceeded {
 			cancel()
 			break
 		}
@@ -300,7 +301,7 @@ func TestGetOnWaitChannel(t *testing.T) {
 	for {
 		var reply string
 		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
-		if err := Invoke(ctx, "/foo/bar", &expectedRequest, &reply, cc, FailFast(false)); Code(err) == codes.DeadlineExceeded {
+		if err := Invoke(ctx, "/foo/bar", &expectedRequest, &reply, cc, FailFast(false)); internal.Code(err) == codes.DeadlineExceeded {
 			cancel()
 			break
 		}
@@ -511,7 +512,7 @@ func TestPickFirstCloseWithPendingRPC(t *testing.T) {
 	// Loop until the above update applies.
 	for {
 		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
-		if err := Invoke(ctx, "/foo/bar", &expectedRequest, &reply, cc, FailFast(false)); Code(err) == codes.DeadlineExceeded {
+		if err := Invoke(ctx, "/foo/bar", &expectedRequest, &reply, cc, FailFast(false)); internal.Code(err) == codes.DeadlineExceeded {
 			cancel()
 			break
 		}

call.go

@@ -29,8 +29,8 @@ import (
 	"google.golang.org/grpc/encoding"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/stats"
-	"google.golang.org/grpc/transport"
 	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/transport"
 )
 
 // recvResponse receives and parses an RPC response.

call_test.go

@@ -31,6 +31,7 @@ import (
 
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/status"
 	"google.golang.org/grpc/test/leakcheck"
 	"google.golang.org/grpc/transport"
@@ -233,7 +234,7 @@ func TestInvokeLargeErr(t *testing.T) {
 	if _, ok := status.FromError(err); !ok {
 		t.Fatalf("grpc.Invoke(_, _, _, _, _) receives non rpc error.")
 	}
-	if Code(err) != codes.Internal || len(errorDesc(err)) != sizeLargeErr {
+	if internal.Code(err) != codes.Internal || len(errorDesc(err)) != sizeLargeErr {
 		t.Fatalf("grpc.Invoke(_, _, _, _, _) = %v, want an error of code %d and desc size %d", err, codes.Internal, sizeLargeErr)
 	}
 	cc.Close()

grpclb/grpclb_test.go

@@ -34,6 +34,7 @@ import (
 	"time"
 
 	"github.com/golang/protobuf/proto"
+	"google.golang.org/grpc/internal"
 
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
@@ -482,7 +483,7 @@ func TestDropRequest(t *testing.T) {
 		for i := 0; i < 3; i++ {
 			// Even RPCs should fail, because the 2st backend has
 			// DropForLoadBalancing set to true.
-			if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(failfast)); grpc.Code(err) != codes.Unavailable {
+			if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(failfast)); internal.Code(err) != codes.Unavailable {
 				t.Errorf("%v.EmptyCall(_, _) = _, %v, want _, %s", testC, err, codes.Unavailable)
 			}
 			// Odd RPCs should succeed since they choose the non-drop-request

internal/internal.go

@@ -19,9 +19,23 @@
 // the godoc of the top-level grpc package.
 package internal
 
+import (
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
 // TestingUseHandlerImpl enables the http.Handler-based server implementation.
 // It must be called before Serve and requires TLS credentials.
 //
 // The provided grpcServer must be of type *grpc.Server. It is untyped
 // for circular dependency reasons.
 var TestingUseHandlerImpl func(grpcServer interface{})
+
+// Code returns the error code for err if it was produced by the rpc system.
+// Otherwise, it returns codes.Unknown.
+func Code(err error) codes.Code {
+	if s, ok := status.FromError(err); ok {
+		return s.Code()
+	}
+	return codes.Unknown
+}

interop/http2/negative_http2_client.go

@@ -34,6 +34,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/interop"
 	testpb "google.golang.org/grpc/interop/grpc_testing"
 )
@@ -77,8 +78,8 @@ func rstAfterHeader(tc testpb.TestServiceClient) {
 	if reply != nil {
 		grpclog.Fatalf("Client received reply despite server sending rst stream after header")
 	}
-	if grpc.Code(err) != codes.Internal {
-		grpclog.Fatalf("%v.UnaryCall() = _, %v, want _, %v", tc, grpc.Code(err), codes.Internal)
+	if internal.Code(err) != codes.Internal {
+		grpclog.Fatalf("%v.UnaryCall() = _, %v, want _, %v", tc, internal.Code(err), codes.Internal)
 	}
 }
 
@@ -88,8 +89,8 @@ func rstDuringData(tc testpb.TestServiceClient) {
 	if reply != nil {
 		grpclog.Fatalf("Client received reply despite server sending rst stream during data")
 	}
-	if grpc.Code(err) != codes.Unknown {
-		grpclog.Fatalf("%v.UnaryCall() = _, %v, want _, %v", tc, grpc.Code(err), codes.Unknown)
+	if internal.Code(err) != codes.Unknown {
+		grpclog.Fatalf("%v.UnaryCall() = _, %v, want _, %v", tc, internal.Code(err), codes.Unknown)
 	}
 }
 
@@ -99,8 +100,8 @@ func rstAfterData(tc testpb.TestServiceClient) {
 	if reply != nil {
 		grpclog.Fatalf("Client received reply despite server sending rst stream after data")
 	}
-	if grpc.Code(err) != codes.Internal {
-		grpclog.Fatalf("%v.UnaryCall() = _, %v, want _, %v", tc, grpc.Code(err), codes.Internal)
+	if internal.Code(err) != codes.Internal {
+		grpclog.Fatalf("%v.UnaryCall() = _, %v, want _, %v", tc, internal.Code(err), codes.Internal)
 	}
 }
 

interop/test_utils.go

@@ -34,6 +34,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal"
 	testpb "google.golang.org/grpc/interop/grpc_testing"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
@@ -231,7 +232,7 @@ func DoTimeoutOnSleepingServer(tc testpb.TestServiceClient, args ...grpc.CallOpt
 	defer cancel()
 	stream, err := tc.FullDuplexCall(ctx, args...)
 	if err != nil {
-		if grpc.Code(err) == codes.DeadlineExceeded {
+		if internal.Code(err) == codes.DeadlineExceeded {
 			return
 		}
 		grpclog.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
@@ -242,11 +243,11 @@ func DoTimeoutOnSleepingServer(tc testpb.TestServiceClient, args ...grpc.CallOpt
 		Payload:      pl,
 	}
 	if err := stream.Send(req); err != nil {
-		if grpc.Code(err) != codes.DeadlineExceeded {
+		if internal.Code(err) != codes.DeadlineExceeded {
 			grpclog.Fatalf("%v.Send(_) = %v", stream, err)
 		}
 	}
-	if _, err := stream.Recv(); grpc.Code(err) != codes.DeadlineExceeded {
+	if _, err := stream.Recv(); internal.Code(err) != codes.DeadlineExceeded {
 		grpclog.Fatalf("%v.Recv() = _, %v, want error code %d", stream, err, codes.DeadlineExceeded)
 	}
 }
@@ -409,8 +410,8 @@ func DoCancelAfterBegin(tc testpb.TestServiceClient, args ...grpc.CallOption) {
 	}
 	cancel()
 	_, err = stream.CloseAndRecv()
-	if grpc.Code(err) != codes.Canceled {
-		grpclog.Fatalf("%v.CloseAndRecv() got error code %d, want %d", stream, grpc.Code(err), codes.Canceled)
+	if internal.Code(err) != codes.Canceled {
+		grpclog.Fatalf("%v.CloseAndRecv() got error code %d, want %d", stream, internal.Code(err), codes.Canceled)
 	}
 }
 
@@ -439,8 +440,8 @@ func DoCancelAfterFirstResponse(tc testpb.TestServiceClient, args ...grpc.CallOp
 		grpclog.Fatalf("%v.Recv() = %v", stream, err)
 	}
 	cancel()
-	if _, err := stream.Recv(); grpc.Code(err) != codes.Canceled {
-		grpclog.Fatalf("%v compleled with error code %d, want %d", stream, grpc.Code(err), codes.Canceled)
+	if _, err := stream.Recv(); internal.Code(err) != codes.Canceled {
+		grpclog.Fatalf("%v compleled with error code %d, want %d", stream, internal.Code(err), codes.Canceled)
 	}
 }
 
@@ -568,15 +569,15 @@ func DoStatusCodeAndMessage(tc testpb.TestServiceClient, args ...grpc.CallOption
 // DoUnimplementedService attempts to call a method from an unimplemented service.
 func DoUnimplementedService(tc testpb.UnimplementedServiceClient) {
 	_, err := tc.UnimplementedCall(context.Background(), &testpb.Empty{})
-	if grpc.Code(err) != codes.Unimplemented {
-		grpclog.Fatalf("%v.UnimplementedCall() = _, %v, want _, %v", tc, grpc.Code(err), codes.Unimplemented)
+	if internal.Code(err) != codes.Unimplemented {
+		grpclog.Fatalf("%v.UnimplementedCall() = _, %v, want _, %v", tc, internal.Code(err), codes.Unimplemented)
 	}
 }
 
 // DoUnimplementedMethod attempts to call an unimplemented method.
 func DoUnimplementedMethod(cc *grpc.ClientConn) {
 	var req, reply proto.Message
-	if err := grpc.Invoke(context.Background(), "/grpc.testing.TestService/UnimplementedCall", req, reply, cc); err == nil || grpc.Code(err) != codes.Unimplemented {
+	if err := grpc.Invoke(context.Background(), "/grpc.testing.TestService/UnimplementedCall", req, reply, cc); err == nil || internal.Code(err) != codes.Unimplemented {
 		grpclog.Fatalf("grpc.Invoke(_, _, _, _, _) = %v, want error code %s", err, codes.Unimplemented)
 	}
 }

pickfirst_test.go

@@ -26,10 +26,11 @@ import (
 
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/resolver/manual"
-	"google.golang.org/grpc/test/leakcheck"
 	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/leakcheck"
 )
 
 func errorDesc(err error) string {
@@ -58,7 +59,7 @@ func TestOneBackendPickfirst(t *testing.T) {
 	defer cancel()
 	req := "port"
 	var reply string
-	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || Code(err) != codes.DeadlineExceeded {
+	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -92,7 +93,7 @@ func TestBackendsPickfirst(t *testing.T) {
 	defer cancel()
 	req := "port"
 	var reply string
-	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || Code(err) != codes.DeadlineExceeded {
+	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -126,7 +127,7 @@ func TestNewAddressWhileBlockingPickfirst(t *testing.T) {
 	defer cancel()
 	req := "port"
 	var reply string
-	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || Code(err) != codes.DeadlineExceeded {
+	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -163,7 +164,7 @@ func TestCloseWithPendingRPCPickfirst(t *testing.T) {
 	defer cancel()
 	req := "port"
 	var reply string
-	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || Code(err) != codes.DeadlineExceeded {
+	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -200,7 +201,7 @@ func TestOneServerDownPickfirst(t *testing.T) {
 	defer cancel()
 	req := "port"
 	var reply string
-	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || Code(err) != codes.DeadlineExceeded {
+	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -242,7 +243,7 @@ func TestAllServersDownPickfirst(t *testing.T) {
 	defer cancel()
 	req := "port"
 	var reply string
-	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || Code(err) != codes.DeadlineExceeded {
+	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}
 
@@ -259,7 +260,7 @@ func TestAllServersDownPickfirst(t *testing.T) {
 		servers[i].stop()
 	}
 	for i := 0; i < 1000; i++ {
-		if err = Invoke(context.Background(), "/foo/bar", &req, &reply, cc); Code(err) == codes.Unavailable {
+		if err = Invoke(context.Background(), "/foo/bar", &req, &reply, cc); internal.Code(err) == codes.Unavailable {
 			return
 		}
 		time.Sleep(time.Millisecond)
@@ -286,7 +287,7 @@ func TestAddressesRemovedPickfirst(t *testing.T) {
 	defer cancel()
 	req := "port"
 	var reply string
-	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || Code(err) != codes.DeadlineExceeded {
+	if err := Invoke(ctx, "/foo/bar", &req, &reply, cc); err == nil || internal.Code(err) != codes.DeadlineExceeded {
 		t.Fatalf("EmptyCall() = _, %v, want _, DeadlineExceeded", err)
 	}