Add custom rooted batch (#60)

ygnmi/ygnmi.go

@@ -22,6 +22,7 @@ import (
 	"reflect"
 	"time"
 
+	"github.com/openconfig/ygot/util"
 	"github.com/openconfig/ygot/ygot"
 	"github.com/openconfig/ygot/ytypes"
 	"google.golang.org/protobuf/encoding/prototext"
@@ -570,3 +571,51 @@ func BatchDelete[T any](sb *SetBatch, q ConfigQuery[T]) {
 		mode: deletePath,
 	})
 }
+
+// Batch contains a collection of paths.
+// Calling State() or Config() on the batch returns a query
+// that can be used to Lookup, Watch, etc on multiple paths at once.
+type Batch[T ygot.ValidatedGoStruct] struct {
+	root  SingletonQuery[T]
+	paths []PathStruct
+}
+
+// NewBatch creates a batch object. All paths in the batch must be children of the root query.
+func NewBatch[T ygot.ValidatedGoStruct](root SingletonQuery[T]) *Batch[T] {
+	return &Batch[T]{
+		root: root,
+	}
+}
+
+// AddPaths adds the paths to the batch. Paths must be children of the root.
+func (b *Batch[T]) AddPaths(paths ...PathStruct) error {
+	root, _, err := ResolvePath(b.root.PathStruct())
+	if err != nil {
+		return err
+	}
+	for _, path := range paths {
+		p, _, err := ResolvePath(path)
+		if err != nil {
+			return err
+		}
+		if !util.PathMatchesQuery(p, root) {
+			return fmt.Errorf("root path %v is not a prefix of %v", root, p)
+		}
+	}
+	b.paths = append(b.paths, paths...)
+	return nil
+}
+
+// Query returns a Query that can be used in gNMI operations.
+// The returned query is immutable, adding paths does not modify existing queries.
+func (b *Batch[T]) Query() SingletonQuery[T] {
+	queryPaths := make([]PathStruct, len(b.paths))
+	copy(queryPaths, b.paths)
+	return NewNonLeafSingletonQuery[T](
+		b.root.dirName(),
+		b.root.IsState(),
+		b.root.PathStruct(),
+		queryPaths,
+		b.root.schema(),
+	)
+}

ygnmi/ygnmi_test.go

@@ -2799,6 +2799,81 @@ func TestBatchWatch(t *testing.T) {
 	}
 }
 
+func TestCustomRootBatch(t *testing.T) {
+	fakeGNMI, c := newClient(t)
+	twoPath := testutil.GNMIPath(t, "/parent/child/state/two")
+
+	tests := []struct {
+		desc                 string
+		stub                 func(s *testutil.Stubber)
+		paths                []ygnmi.PathStruct
+		wantSubscriptionPath []*gpb.Path
+		wantVal              *ygnmi.Value[*exampleoc.Parent]
+		wantAddErr           string
+		wantLookupErr        string
+	}{{
+		desc: "not prefix",
+		stub: func(s *testutil.Stubber) {},
+		paths: []ygnmi.PathStruct{
+			exampleocpath.Root().Model(),
+		},
+		wantAddErr: "is not a prefix",
+	}, {
+		desc: "success",
+		stub: func(s *testutil.Stubber) {
+			s.Notification(&gpb.Notification{
+				Timestamp: 100,
+				Update: []*gpb.Update{{
+					Path: twoPath,
+					Val:  &gpb.TypedValue{Value: &gpb.TypedValue_StringVal{StringVal: "foo"}},
+				}},
+			}).Sync()
+		},
+		paths: []ygnmi.PathStruct{
+			exampleocpath.Root().Parent().Child().Two(),
+		},
+		wantSubscriptionPath: []*gpb.Path{
+			twoPath,
+		},
+		wantVal: (&ygnmi.Value[*exampleoc.Parent]{
+			Timestamp: time.Unix(0, 100),
+			Path:      testutil.GNMIPath(t, "/parent"),
+		}).SetVal(&exampleoc.Parent{
+			Child: &exampleoc.Parent_Child{
+				Two: ygot.String("foo"),
+			},
+		}),
+	}}
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			tt.stub(fakeGNMI.Stub())
+			b := ygnmi.NewBatch(exampleocpath.Root().Parent().State())
+			err := b.AddPaths(tt.paths...)
+			if diff := errdiff.Substring(err, tt.wantAddErr); diff != "" {
+				t.Fatalf("AddPaths returned unexpected diff: %s", diff)
+			}
+			if err != nil {
+				return
+			}
+			got, gotErr := ygnmi.Lookup(context.Background(), c, b.Query())
+			if diff := errdiff.Substring(gotErr, tt.wantLookupErr); diff != "" {
+				t.Fatalf("Watch() returned unexpected diff: %s", diff)
+			}
+			if gotErr != nil {
+				return
+			}
+			checkJustReceived(t, got.RecvTimestamp)
+			verifySubscriptionPathsSent(t, fakeGNMI, tt.wantSubscriptionPath...)
+			tt.wantVal.RecvTimestamp = got.RecvTimestamp
+
+			if diff := cmp.Diff(tt.wantVal, got, cmp.AllowUnexported(ygnmi.Value[*exampleoc.Parent]{}), protocmp.Transform()); diff != "" {
+				t.Errorf("Watch() returned unexpected diff (-want,+got):\n %s\nComplianceErrors:\n%v", diff, got.ComplianceErrors)
+			}
+		})
+	}
+
+}
+
 func TestSetBatch(t *testing.T) {
 	setClient := &fakeGNMISetClient{}
 	client, err := ygnmi.NewClient(setClient, ygnmi.WithTarget("dut"))