Project import generated by Copybara.

FolderOrigin-RevId: /usr/local/google/home/gdennis/copybara/temp/folder-destination10187328992434430653/.

binding/grpcutil/grpcutil.go

@@ -0,0 +1,100 @@
+// Copyright 2022 Google LLC
+//
+// 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
+//
+//     https://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 grpcutil contains gRPC utilities useful to binding implementations.
+package grpcutil
+
+import (
+	"golang.org/x/net/context"
+	"io"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+var emptyCancel context.CancelFunc = func() {}
+
+// WithDefaultTimeout returns context.WithTimeout(ctx, timeout) if ctx has no
+// deadline; otherwise returns the context as-is and an empty cancel function.
+func WithDefaultTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
+	if _, ok := ctx.Deadline(); !ok {
+		return context.WithTimeout(ctx, timeout)
+	}
+	return ctx, emptyCancel
+}
+
+// WithUnaryDefaultTimeout returns a dial option that intercepts unary requests
+// and imposes the specified context timeout on the request if it does not
+// already have a deadline.
+func WithUnaryDefaultTimeout(timeout time.Duration) grpc.DialOption {
+	return grpc.WithChainUnaryInterceptor(
+		func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+			ctx, cancel := WithDefaultTimeout(ctx, timeout)
+			defer cancel()
+			return invoker(ctx, method, req, reply, cc, opts...)
+		},
+	)
+}
+
+// WithStreamDefaultTimeout returns a dial option that intercepts stream
+// requests and imposes the specified context timeout on the request if it does
+// not already have a deadline.
+func WithStreamDefaultTimeout(timeout time.Duration) grpc.DialOption {
+	return grpc.WithChainStreamInterceptor(
+		func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+			ctx, cancel := WithDefaultTimeout(ctx, timeout)
+			cs, err := streamer(ctx, desc, cc, method, opts...)
+			if err != nil {
+				cancel()
+				return nil, err
+			}
+			// Unfortunately gRPC provides no simple way to ensure a child context
+			// created in an interceptor is cancelled.
+			// The current recommendation is to wrap the client stream and cancel it
+			// when RecvMsg returns a non-nil error, or when Header or SendMsg return
+			// a non-nil, non-EOF error. See https://github.com/grpc/grpc-go/issues/1717
+			return &cancelClientStream{ClientStream: cs, cancel: cancel}, nil
+		},
+	)
+}
+
+type cancelClientStream struct {
+	grpc.ClientStream
+	cancel context.CancelFunc
+}
+
+func (cs *cancelClientStream) Header() (metadata.MD, error) {
+	md, err := cs.ClientStream.Header()
+	if err != nil && err != io.EOF {
+		cs.cancel()
+	}
+	return md, err
+}
+
+func (cs *cancelClientStream) SendMsg(m any) error {
+	err := cs.ClientStream.SendMsg(m)
+	if err != nil && err != io.EOF {
+		cs.cancel()
+	}
+	return err
+}
+
+func (cs *cancelClientStream) RecvMsg(m any) error {
+	err := cs.ClientStream.RecvMsg(m)
+	if err != nil {
+		cs.cancel()
+	}
+	return err
+}

binding/grpcutil/grpcutil_test.go

@@ -0,0 +1,106 @@
+// Copyright 2022 Google LLC
+//
+// 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
+//
+//     https://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 grpcutil
+
+import (
+	"golang.org/x/net/context"
+	"testing"
+	"time"
+
+	"github.com/openconfig/ondatra/binding/grpcutil/testservice"
+
+	tgrpcpb "github.com/openconfig/ondatra/binding/grpcutil/testservice/gen"
+	tpb "github.com/openconfig/ondatra/binding/grpcutil/testservice/gen"
+)
+
+func TestDefaultTimeout(t *testing.T) {
+	const (
+		defaultTimeout = time.Minute
+		givenTimeout   = 10 * defaultTimeout
+	)
+	spy := new(spyServer)
+	ctx := context.Background()
+	serv := testservice.Start(ctx, t, spy,
+		WithUnaryDefaultTimeout(defaultTimeout),
+		WithStreamDefaultTimeout(defaultTimeout))
+	defer serv.Stop()
+
+	t.Run("unary default deadline", func(t *testing.T) {
+		spy.deadline = time.Time{}
+		wantDeadline := time.Now().Add(defaultTimeout)
+		serv.MustSendUnary(ctx, t, "")
+		if gotDeadline := spy.deadline; !timeApproxEq(gotDeadline, wantDeadline) {
+			t.Errorf("WithUnaryDefaultTimeout() got deadline %v, want approximately %v", gotDeadline, wantDeadline)
+		}
+	})
+
+	t.Run("unary given deadline", func(t *testing.T) {
+		spy.deadline = time.Time{}
+		wantDeadline := time.Now().Add(givenTimeout)
+		ctx, cancel := context.WithTimeout(ctx, givenTimeout)
+		defer cancel()
+		serv.MustSendUnary(ctx, t, "")
+		if gotDeadline := spy.deadline; !timeApproxEq(gotDeadline, wantDeadline) {
+			t.Errorf("WithUnaryDefaultTimeout() got deadline %v, want approximately %v", gotDeadline, wantDeadline)
+		}
+	})
+
+	t.Run("stream default deadline", func(t *testing.T) {
+		spy.deadline = time.Time{}
+		wantDeadline := time.Now().Add(defaultTimeout)
+		streamClient := serv.MustSendStream(ctx, t)
+		for streamClient.MustRecv(t) {
+		}
+		if gotDeadline := spy.deadline; !timeApproxEq(gotDeadline, wantDeadline) {
+			t.Errorf("WithStreamDefaultTimeout() got deadline %v, want approximately %v", gotDeadline, wantDeadline)
+		}
+	})
+
+	t.Run("stream given deadline", func(t *testing.T) {
+		spy.deadline = time.Time{}
+		wantDeadline := time.Now().Add(givenTimeout)
+		ctx, cancel := context.WithTimeout(ctx, givenTimeout)
+		defer cancel()
+		streamClient := serv.MustSendStream(ctx, t)
+		for streamClient.MustRecv(t) {
+		}
+		if gotDeadline := spy.deadline; !timeApproxEq(gotDeadline, wantDeadline) {
+			t.Errorf("WithStreamDefaultTimeout() got deadline %v, want approximately %v", gotDeadline, wantDeadline)
+		}
+	})
+}
+
+type spyServer struct {
+	tgrpcpb.UnimplementedTestServer
+	deadline time.Time
+}
+
+func (s *spyServer) SendUnary(ctx context.Context, _ *tpb.TestRequest) (*tpb.TestResponse, error) {
+	if deadline, ok := ctx.Deadline(); ok {
+		s.deadline = deadline
+	}
+	return new(tpb.TestResponse), nil
+}
+
+func (s *spyServer) SendStream(stream tgrpcpb.Test_SendStreamServer) error {
+	if deadline, ok := stream.Context().Deadline(); ok {
+		s.deadline = deadline
+	}
+	return nil
+}
+
+func timeApproxEq(t1, t2 time.Time) bool {
+	return t1.Sub(t2).Abs().Seconds() < 1
+}

internal/rawapis/testservice/generate.sh → binding/grpcutil/testservice/generate.sh

@@ -19,5 +19,6 @@
 set -e
 
 cd "$( dirname "${BASH_SOURCE[0]}" )"
-protoc --go_out=. --go_opt=paths=source_relative --go-grpc_out=. \
+mkdir -p gen
+protoc --go_out=gen --go_opt=paths=source_relative --go-grpc_out=gen \
   --go-grpc_opt=paths=source_relative testservice.proto

binding/grpcutil/testservice/testservice.go

@@ -0,0 +1,104 @@
+// Copyright 2022 Google LLC
+//
+// 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
+//
+//     https://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 testservice provides an interface for testing with the Test service.
+package testservice
+
+import (
+	"golang.org/x/net/context"
+	"io"
+	"net"
+	"testing"
+
+	"google.golang.org/grpc/credentials/local"
+	"google.golang.org/grpc"
+
+	tgrpcpb "github.com/openconfig/ondatra/binding/grpcutil/testservice/gen"
+	tpb "github.com/openconfig/ondatra/binding/grpcutil/testservice/gen"
+)
+
+// Start starts the gRCP server with the specified TestServer implementation,
+// dials the server, and returns a handle to the test service.
+func Start(ctx context.Context, t *testing.T, testSrv tgrpcpb.TestServer, opts ...grpc.DialOption) *TestService {
+	t.Helper()
+	srv := grpc.NewServer()
+	tgrpcpb.RegisterTestServer(srv, testSrv)
+	lis, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("net.Listen() failed: %v", err)
+	}
+	go srv.Serve(lis)
+
+	opts = append([]grpc.DialOption{grpc.WithTransportCredentials(local.NewCredentials())}, opts...)
+	conn, err := grpc.DialContext(ctx, lis.Addr().String(), opts...)
+	if err != nil {
+		t.Fatalf("grpc.DialContext() failed: %v", err)
+	}
+	return &TestService{TestClient: tgrpcpb.NewTestClient(conn), stopFn: srv.Stop}
+}
+
+// TestService represents a running test service client and server.
+type TestService struct {
+	tgrpcpb.TestClient
+	stopFn func()
+}
+
+// MustSendUnary calls SendUnary and fails the test fatally on error.
+func (ts *TestService) MustSendUnary(ctx context.Context, t *testing.T, msg string) {
+	t.Helper()
+	if _, err := ts.SendUnary(ctx, &tpb.TestRequest{Message: msg}); err != nil {
+		t.Fatalf("SendUnary() failed: %v", err)
+	}
+}
+
+// MustSendStream calls SendStream and fails the test fatally on error.
+func (ts *TestService) MustSendStream(ctx context.Context, t *testing.T) *TestStreamClient {
+	t.Helper()
+	streamClient, err := ts.SendStream(ctx)
+	if err != nil {
+		t.Fatalf("SendStream() failed: %v", err)
+	}
+	return &TestStreamClient{streamClient}
+}
+
+// Stop stops the TestService server.
+func (ts *TestService) Stop() {
+	ts.stopFn()
+}
+
+// TestStreamClient is a SendStreamClient with additional Must methods.
+type TestStreamClient struct {
+	tgrpcpb.Test_SendStreamClient
+}
+
+// MustSend calls Send and fails the test fatally on error.
+func (sc *TestStreamClient) MustSend(t *testing.T, msg string) {
+	t.Helper()
+	if err := sc.Send(&tpb.TestRequest{Message: msg}); err != nil {
+		t.Fatalf("Send() failed: %v", err)
+	}
+}
+
+// MustRecv calls Recv and fails the test fatally on a non-EOF error.
+// Returns true if a message is received (if there is no error or an EOF error).
+func (sc *TestStreamClient) MustRecv(t *testing.T) bool {
+	t.Helper()
+	if _, err := sc.Recv(); err != nil {
+		if err == io.EOF {
+			return false
+		}
+		t.Fatalf("Recv() failed: %v", err)
+	}
+	return true
+}

internal/rawapis/testservice/testservice.proto → binding/grpcutil/testservice/testservice.proto

@@ -16,7 +16,7 @@ syntax = "proto3";
 
 package testservice;
 
-option go_package = "github.com/openconfig/ondatra/internal/rawapis/testservice";
+option go_package = "github.com/openconfig/ondatra/binding/grpcutil/testservice";
 
 service Test {
   // TestUnary is a fake unary RPC.

config/acl/acl.go

@@ -3,7 +3,7 @@ Package acl is a generated package which contains definitions
 of structs which generate gNMI paths for a YANG schema. The generated paths are
 based on a compressed form of the schema.
 
-This package was generated by /usr/local/google/home/alexmasi/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
+This package was generated by /usr/local/google/home/gdennis/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
 using the following YANG input files:
   - gnmi-collector-metadata.yang
   - gnsi/authz/gnsi-authz.yang

config/device/device.go

@@ -3,7 +3,7 @@ Package device is a generated package which contains definitions
 of structs which generate gNMI paths for a YANG schema. The generated paths are
 based on a compressed form of the schema.
 
-This package was generated by /usr/local/google/home/alexmasi/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
+This package was generated by /usr/local/google/home/gdennis/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
 using the following YANG input files:
   - gnmi-collector-metadata.yang
   - gnsi/authz/gnsi-authz.yang

config/gnmicollectormetadata/gnmicollectormetadata.go

@@ -3,7 +3,7 @@ Package gnmicollectormetadata is a generated package which contains definitions
 of structs which generate gNMI paths for a YANG schema. The generated paths are
 based on a compressed form of the schema.
 
-This package was generated by /usr/local/google/home/alexmasi/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
+This package was generated by /usr/local/google/home/gdennis/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
 using the following YANG input files:
   - gnmi-collector-metadata.yang
   - gnsi/authz/gnsi-authz.yang

config/interfaces/interfaces.go

@@ -3,7 +3,7 @@ Package interfaces is a generated package which contains definitions
 of structs which generate gNMI paths for a YANG schema. The generated paths are
 based on a compressed form of the schema.
 
-This package was generated by /usr/local/google/home/alexmasi/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
+This package was generated by /usr/local/google/home/gdennis/go/pkg/mod/github.com/openconfig/[email protected]/genutil/names.go
 using the following YANG input files:
   - gnmi-collector-metadata.yang
   - gnsi/authz/gnsi-authz.yang