Skip to content

Commit

Permalink
Project import generated by Copybara.
Browse files Browse the repository at this point in the history 
FolderOrigin-RevId: /usr/local/google/home/gdennis/copybara/temp/folder-destination11220037195300927884/.
  • Loading branch information
GGN Engprod Team authored and greg-dennis committed Nov 8, 2022
1 parent 3f0b6dd commit 0dbcef7
Show file tree
Hide file tree
Showing 17 changed files with 84,937 additions and 84,895 deletions.
8 changes: 5 additions & 3 deletions device.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ import (
opb "github.com/openconfig/ondatra/proto"
)

var _ gnmi.DeviceOrOpts = (*Device)(nil)

// Device represents a network device.
type Device struct {
id string
Expand Down Expand Up @@ -81,10 +83,10 @@ func (v Vendor) String() string {
return opb.Device_Vendor(v).String()
}

// GNMI returns a gNMI client for the device.
func (d *Device) GNMI() *gnmi.Client {
// GNMIOpts returns a new set of options to customize gNMI queries.
func (d *Device) GNMIOpts() *gnmi.Opts {
useGetForCfg := d.Vendor() == CISCO || d.Vendor() == JUNIPER
return gnmi.NewClient(d.Name(), useGetForCfg, d.clientFn)
return gnmi.NewOpts(d.Name(), useGetForCfg, d.clientFn)
}

// Telemetry returns a telemetry path root for the device.
Expand Down
127 changes: 76 additions & 51 deletions gnmi/gnmi.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ import (

"golang.org/x/net/context"

"github.com/openconfig/ondatra/gnmi/oc/ocpath"
"github.com/openconfig/ondatra/gnmi/otg/otgpath"
"github.com/openconfig/ygnmi/ygnmi"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
Expand All @@ -30,8 +32,30 @@ import (
gpb "github.com/openconfig/gnmi/proto/gnmi"
)

// Client is gnmi client that supports custom options.
type Client struct {
// OC returns the root OpenConfig gNMI path.
func OC() *ocpath.RootPath {
return ocpath.Root()
}

// OCBatch returns a new batch collection of OpenConfig gnmi paths.
func OCBatch() *ocpath.Batch {
return new(ocpath.Batch)
}

// OTG returns the root Open Traffic Generator gNMI path.
func OTG() *otgpath.RootPath {
return otgpath.Root()
}

// OTGBatch returns a new batch collection of Open Traffic Generator gnmi paths.
func OTGBatch() *otgpath.Batch {
return new(otgpath.Batch)
}

var _ DeviceOrOpts = (*Opts)(nil)

// Opts contains customs options for a gNMI request.
type Opts struct {
id string
initClientFn func(context.Context) (gpb.GNMIClient, error)
md metadata.MD
Expand All @@ -40,43 +64,44 @@ type Client struct {
opts []ygnmi.Option
}

// NewClient creates a new gNMI client. DO NOT USE: call dut.GNMI() or ate.OTG().GNMI() instead.
func NewClient(id string, useGetForConfig bool, initClient func(context.Context) (gpb.GNMIClient, error)) *Client {
return &Client{
// NewOpts creates a set of GNMI options.
// DO NOT USE: call dut.GNMIOpts() or ate.OTG().GNMIOpts() instead.
func NewOpts(id string, useGetForConfig bool, initClient func(context.Context) (gpb.GNMIClient, error)) *Opts {
return &Opts{
initClientFn: initClient,
id: id,
useGetForConfig: useGetForConfig,
}
}

// WithMetadata adds metadata to the outgoing context, overriding any set by previous calls.
func (c *Client) WithMetadata(md metadata.MD) *Client {
c.md = md
return c
func (o *Opts) WithMetadata(md metadata.MD) *Opts {
o.md = md
return o
}

// WithClient sets a custom gNMI client to the client, overriding any set by previous calls.
func (c *Client) WithClient(client gpb.GNMIClient) *Client {
c.client = client
return c
func (o *Opts) WithClient(client gpb.GNMIClient) *Opts {
o.client = client
return o
}

// WithYGNMIOpts adds the ygnmi Options to the client, overriding any set by previous calls.
func (c *Client) WithYGNMIOpts(opts ...ygnmi.Option) *Client {
c.opts = opts
return c
func (o *Opts) WithYGNMIOpts(opts ...ygnmi.Option) *Opts {
o.opts = opts
return o
}

// GNMI implements the DUTOrGNMI interface.
func (c *Client) GNMI() *Client { return c }
// GNMIOpts implements the DeviceOrOpts interface.
func (o *Opts) GNMIOpts() *Opts { return o }

// DeviceOrGNMI is an interface that is ondatra.DUT or a gnmi.Client
type DeviceOrGNMI interface {
GNMI() *Client
// DeviceOrOpts is an interface that is ondatra.Device or a gnmi.Opts
type DeviceOrOpts interface {
GNMIOpts() *Opts
}

// Lookup fetches the value of a ygnmi.SingletonQuery with a ONCE subscription.
func Lookup[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T]) *ygnmi.Value[T] {
func Lookup[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.SingletonQuery[T]) *ygnmi.Value[T] {
t.Helper()
c := newClient(t, dev, "Lookup")
v, err := ygnmi.Lookup(createContext(dev), c, q, createGetOpts[T](dev, q)...)
Expand All @@ -89,14 +114,14 @@ func Lookup[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T]) *y
// LookupConfig fetches the value of a ygnmi.SingletonQuery with a ONCE subscription.
// Note: This is a workaround for Go's type inference not working for this use case and may be removed in a subsequent release.
// Note: This is equivalent to calling Lookup with a ConfigQuery and providing a fully-qualified type parameter.
func LookupConfig[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.ConfigQuery[T]) *ygnmi.Value[T] {
func LookupConfig[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.ConfigQuery[T]) *ygnmi.Value[T] {
return Lookup[T](t, dev, q)
}

// Get fetches the value of a SingletonQuery with a ONCE subscription,
// failing the test fatally if no value is present at the path.
// Use Lookup to also get metadata or tolerate no value present.
func Get[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T]) T {
func Get[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.SingletonQuery[T]) T {
t.Helper()
c := newClient(t, dev, "Get")
v, err := ygnmi.Get(createContext(dev), c, q, createGetOpts[T](dev, q)...)
Expand All @@ -111,7 +136,7 @@ func Get[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T]) T {
// Use Lookup to also get metadata or tolerate no value present.
// Note: This is a workaround for Go's type inference not working for this use case and may be removed in a subsequent release.
// Note: This is equivalent to calling Get with a ConfigQuery and providing a fully-qualified type parameter.
func GetConfig[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.ConfigQuery[T]) T {
func GetConfig[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.ConfigQuery[T]) T {
return Get[T](t, dev, q)
}

Expand All @@ -123,7 +148,7 @@ type watchAwaiter[T any] interface {
type Watcher[T any] struct {
watcher watchAwaiter[T]
cancelFn func()
dev DeviceOrGNMI
dev DeviceOrOpts
query ygnmi.AnyQuery[T]
}

Expand Down Expand Up @@ -167,15 +192,15 @@ func (w *Watcher[T]) Cancel() {
// or the timeout is reached. Calling Await on the returned Watcher waits for the subscription
// to complete. It returns the last observed value and a boolean that indicates whether
// that value satisfies the predicate.
func Watch[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T], timeout time.Duration, pred func(*ygnmi.Value[T]) bool) *Watcher[T] {
func Watch[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.SingletonQuery[T], timeout time.Duration, pred func(*ygnmi.Value[T]) bool) *Watcher[T] {
c := newClient(t, dev, "Watch")
ctx, cancel := context.WithTimeout(createContext(dev), timeout)
w := ygnmi.Watch(ctx, c, q, func(v *ygnmi.Value[T]) error {
if ok := pred(v); ok {
return nil
}
return ygnmi.Continue
}, dev.GNMI().opts...)
}, dev.GNMIOpts().opts...)
return &Watcher[T]{
watcher: w,
cancelFn: cancel,
Expand All @@ -188,13 +213,13 @@ func Watch[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T], tim
// blocking until a value that is deep equal to the specified val is received
// or the timeout is reached. To wait for a generic predicate, or to make a
// non-blocking call, use the Watch method instead.
func Await[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T], timeout time.Duration, val T) *ygnmi.Value[T] {
func Await[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.SingletonQuery[T], timeout time.Duration, val T) *ygnmi.Value[T] {
t.Helper()
c := newClient(t, dev, "Await")
ctx, cancel := context.WithTimeout(createContext(dev), timeout)
defer cancel()

v, err := ygnmi.Await(ctx, c, q, val, dev.GNMI().opts...)
v, err := ygnmi.Await(ctx, c, q, val, dev.GNMIOpts().opts...)
if err != nil {
t.Fatalf("Await(t) on %s at %v: %v", dev, q, err)
}
Expand All @@ -209,7 +234,7 @@ type collectAwaiter[T any] interface {
type Collector[T any] struct {
collector collectAwaiter[T]
cancelFn func()
dev DeviceOrGNMI
dev DeviceOrOpts
query ygnmi.AnyQuery[T]
}

Expand All @@ -230,11 +255,11 @@ func (c *Collector[T]) Await(t testing.TB) []*ygnmi.Value[T] {

// Collect starts an asynchronous collection of the values at the query with a STREAM subscription.
// Calling Await on the return Collection waits until the timeout is reached and returns the collected values.
func Collect[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T], timeout time.Duration) *Collector[T] {
func Collect[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.SingletonQuery[T], timeout time.Duration) *Collector[T] {
c := newClient(t, dev, "Collect")
ctx, cancel := context.WithTimeout(createContext(dev), timeout)
collect := &Collector[T]{
collector: ygnmi.Collect(ctx, c, q, dev.GNMI().opts...),
collector: ygnmi.Collect(ctx, c, q, dev.GNMIOpts().opts...),
cancelFn: cancel,
dev: dev,
query: q,
Expand All @@ -245,7 +270,7 @@ func Collect[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.SingletonQuery[T], t

// LookupAll fetches the values of a ygnmi.WildcardQuery with a ONCE subscription.
// It returns an empty list if no values are present at the path.
func LookupAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T]) []*ygnmi.Value[T] {
func LookupAll[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.WildcardQuery[T]) []*ygnmi.Value[T] {
t.Helper()
c := newClient(t, dev, "LookupAll")
v, err := ygnmi.LookupAll(createContext(dev), c, q, createGetOpts[T](dev, q)...)
Expand All @@ -258,7 +283,7 @@ func LookupAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T])
// GetAll fetches the value of a WildcardQuery with a ONCE subscription skipping any non-present paths.
// It fails the test fatally if no value is present at the path
// Use LookupAll to also get metadata or tolerate no values present.
func GetAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T]) []T {
func GetAll[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.WildcardQuery[T]) []T {
t.Helper()
c := newClient(t, dev, "GetAll")
v, err := ygnmi.GetAll(createContext(dev), c, q, createGetOpts[T](dev, q)...)
Expand All @@ -273,15 +298,15 @@ func GetAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T]) []T
// or the timeout is reached. Calling Await on the returned Watcher waits for the subscription
// to complete. It returns the last observed value and a boolean that indicates whether
// that value satisfies the predicate.
func WatchAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T], timeout time.Duration, pred func(*ygnmi.Value[T]) bool) *Watcher[T] {
func WatchAll[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.WildcardQuery[T], timeout time.Duration, pred func(*ygnmi.Value[T]) bool) *Watcher[T] {
c := newClient(t, dev, "WatchAll")
ctx, cancel := context.WithTimeout(createContext(dev), timeout)
w := ygnmi.WatchAll(ctx, c, q, func(v *ygnmi.Value[T]) error {
if ok := pred(v); ok {
return nil
}
return ygnmi.Continue
}, dev.GNMI().opts...)
}, dev.GNMIOpts().opts...)
return &Watcher[T]{
watcher: w,
cancelFn: cancel,
Expand All @@ -292,11 +317,11 @@ func WatchAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T], t

// CollectAll starts an asynchronous collection of the values at the query with a STREAM subscription.
// Calling Await on the return Collection waits until the timeout is reached and returns the collected values.
func CollectAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T], timeout time.Duration) *Collector[T] {
func CollectAll[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.WildcardQuery[T], timeout time.Duration) *Collector[T] {
c := newClient(t, dev, "CollectAll")
ctx, cancel := context.WithTimeout(createContext(dev), timeout)
collect := &Collector[T]{
collector: ygnmi.CollectAll(ctx, c, q, dev.GNMI().opts...),
collector: ygnmi.CollectAll(ctx, c, q, dev.GNMIOpts().opts...),
cancelFn: cancel,
dev: dev,
query: q,
Expand All @@ -306,7 +331,7 @@ func CollectAll[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.WildcardQuery[T],
}

// Update updates the configuration at the given query path with the val.
func Update[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.ConfigQuery[T], val T) *ygnmi.Result {
func Update[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.ConfigQuery[T], val T) *ygnmi.Result {
t.Helper()
c := newClient(t, dev, "Update")
res, err := ygnmi.Update(createContext(dev), c, q, val)
Expand All @@ -317,7 +342,7 @@ func Update[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.ConfigQuery[T], val T
}

// Replace replaces the configuration at the given query path with the val.
func Replace[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.ConfigQuery[T], val T) *ygnmi.Result {
func Replace[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.ConfigQuery[T], val T) *ygnmi.Result {
t.Helper()
c := newClient(t, dev, "Replace")
res, err := ygnmi.Replace(createContext(dev), c, q, val)
Expand All @@ -328,7 +353,7 @@ func Replace[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.ConfigQuery[T], val
}

// Delete deletes the configuration at the given query path.
func Delete[T any](t testing.TB, dev DeviceOrGNMI, q ygnmi.ConfigQuery[T]) *ygnmi.Result {
func Delete[T any](t testing.TB, dev DeviceOrOpts, q ygnmi.ConfigQuery[T]) *ygnmi.Result {
t.Helper()
c := newClient(t, dev, "Delete")
res, err := ygnmi.Delete(createContext(dev), c, q)
Expand All @@ -345,7 +370,7 @@ type SetBatch struct {
}

// Set performs the gnmi.Set request with all queued operations.
func (sb *SetBatch) Set(t testing.TB, dev DeviceOrGNMI) *ygnmi.Result {
func (sb *SetBatch) Set(t testing.TB, dev DeviceOrOpts) *ygnmi.Result {
t.Helper()
c := newClient(t, dev, "Set")
res, err := sb.sb.Set(createContext(dev), c)
Expand All @@ -370,34 +395,34 @@ func BatchDelete[T any](sb *SetBatch, q ygnmi.ConfigQuery[T]) {
ygnmi.BatchDelete(sb.sb, q)
}

func createContext(d DeviceOrGNMI) context.Context {
func createContext(d DeviceOrOpts) context.Context {
ctx := context.Background()
if d.GNMI().md == nil {
if d.GNMIOpts().md == nil {
return ctx
}
return metadata.NewOutgoingContext(ctx, d.GNMI().md)
return metadata.NewOutgoingContext(ctx, d.GNMIOpts().md)
}

func createGetOpts[T any](d DeviceOrGNMI, q ygnmi.AnyQuery[T]) []ygnmi.Option {
opts := d.GNMI().opts
if d.GNMI().useGetForConfig && !q.IsState() {
func createGetOpts[T any](d DeviceOrOpts, q ygnmi.AnyQuery[T]) []ygnmi.Option {
opts := d.GNMIOpts().opts
if d.GNMIOpts().useGetForConfig && !q.IsState() {
opts = append(opts, ygnmi.WithUseGet())
}
return opts
}

func newClient(t testing.TB, dev DeviceOrGNMI, method string) *ygnmi.Client {
func newClient(t testing.TB, dev DeviceOrOpts, method string) *ygnmi.Client {
t.Helper()
gnmiC := dev.GNMI().client
gnmiC := dev.GNMIOpts().client
if gnmiC == nil {
var err error
gnmiC, err = dev.GNMI().initClientFn(context.Background())
gnmiC, err = dev.GNMIOpts().initClientFn(context.Background())
if err != nil {
t.Fatalf("Failed to get client: %v", err)
}
}

c, err := ygnmi.NewClient(gnmiC, ygnmi.WithTarget(dev.GNMI().id))
c, err := ygnmi.NewClient(gnmiC, ygnmi.WithTarget(dev.GNMIOpts().id))
if err != nil {
t.Fatalf("%s(t) on %s: %v", method, dev, err)
}
Expand Down
6 changes: 2 additions & 4 deletions gnmi/migrator/gnmifix/gnmifix.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,6 @@ const (
newImports = `
"github.com/openconfig/ygnmi/ygnmi"
"github.com/openconfig/ondatra/gnmi"
"github.com/openconfig/ondatra/gnmi/oc/ocpath"
"github.com/openconfig/ondatra/gnmi/otg/otgpath"
"github.com/openconfig/ondatra/gnmi/oc"
`
oldOTGTypesImport = `"github.com/openconfig/ondatra/telemetry/otg"`
Expand Down Expand Up @@ -210,9 +208,9 @@ type callInfo struct {
func (c callInfo) pathToString() string {
left := strings.TrimSuffix(c.caller, ".")
if c.pathFound && !c.indirect {
left = "ocpath.Root()"
left = "gnmi.OC()"
if c.otg {
left = "otgpath.Root()"
left = "gnmi.OTG()"
}
}

Expand Down
Loading

0 comments on commit 0dbcef7

Please sign in to comment.