Add FakeRIB for testing.

rib/helpers.go

@@ -17,14 +17,16 @@ package rib
 import (
 	"fmt"
 
+	"github.com/openconfig/ygot/ygot"
+
 	aftpb "github.com/openconfig/gribi/v1/proto/gribi_aft"
 	spb "github.com/openconfig/gribi/v1/proto/service"
-	"github.com/openconfig/ygot/ygot"
+	wpb "github.com/openconfig/ygot/proto/ywrapper"
 )
 
-// RIBFromGetResponses returns a RIB from a slice of gRIBI GetResponse messages.
+// FromGetResponses returns a RIB from a slice of gRIBI GetResponse messages.
 // The supplied defaultName is used as the default network instance name.
-func RIBFromGetResponses(defaultName string, responses []*spb.GetResponse) (*RIB, error) {
+func FromGetResponses(defaultName string, responses []*spb.GetResponse) (*RIB, error) {
 	r := New(defaultName)
 	niAFTs := map[string]*aftpb.Afts{}
 
@@ -70,3 +72,89 @@ func RIBFromGetResponses(defaultName string, responses []*spb.GetResponse) (*RIB
 
 	return r, nil
 }
+
+// fakeRIB is a RIB for use in testing which exposes methods that can be used to more easily
+// construct a RIB's contents.
+type fakeRIB struct {
+	r *RIB
+}
+
+// NewFake returns a new Fake RIB.
+func NewFake(defaultName string, opt ...RIBOpt) *fakeRIB {
+	return &fakeRIB{
+		r: New(defaultName, opt...),
+	}
+}
+
+// RIB returns the constructed fake RIB to the caller.
+func (f *fakeRIB) RIB() *RIB {
+	return f.r
+}
+
+// InjectIPv4 adds an IPv4 entry to network instance ni, with the specified
+// prefix (pfx), and referencing the specified next-hop-group with index nhg.
+// It returns an error if the entry cannot be injected.
+func (f *fakeRIB) InjectIPv4(ni, pfx string, nhg uint64) error {
+	niR, ok := f.r.NetworkInstanceRIB(ni)
+	if !ok {
+		return fmt.Errorf("unknown NI, %s", ni)
+	}
+	if _, _, err := niR.AddIPv4(&aftpb.Afts_Ipv4EntryKey{
+		Prefix: pfx,
+		Ipv4Entry: &aftpb.Afts_Ipv4Entry{
+			NextHopGroup: &wpb.UintValue{Value: nhg},
+		},
+	}, false); err != nil {
+		return fmt.Errorf("cannot add IPv4 entry, err: %v", err)
+	}
+
+	return nil
+}
+
+// InjectNHG adds a next-hop-group entry to network instance ni, with the specified
+// ID (nhgId). The next-hop-group contains the next hops specified in the nhs map,
+// with the key of the map being the next-hop ID and the value being the weight within
+// the group.
+func (f *fakeRIB) InjectNHG(ni string, nhgId uint64, nhs map[uint64]uint64) error {
+	niR, ok := f.r.NetworkInstanceRIB(ni)
+	if !ok {
+		return fmt.Errorf("unknown NI, %s", ni)
+	}
+
+	nhg := &aftpb.Afts_NextHopGroupKey{
+		Id:           nhgId,
+		NextHopGroup: &aftpb.Afts_NextHopGroup{},
+	}
+	for nh, weight := range nhs {
+		nhg.NextHopGroup.NextHop = append(nhg.NextHopGroup.NextHop, &aftpb.Afts_NextHopGroup_NextHopKey{
+			Index: nh,
+			NextHop: &aftpb.Afts_NextHopGroup_NextHop{
+				Weight: &wpb.UintValue{Value: weight},
+			},
+		})
+	}
+
+	if _, _, err := niR.AddNextHopGroup(nhg, false); err != nil {
+		return fmt.Errorf("cannot add NHG entry, err: %v", err)
+	}
+
+	return nil
+}
+
+// InjectNH adds a next-hop entry to network instance ni, with the specified
+// index (nhIdx). An error is returned if it cannot be added.
+func (f *fakeRIB) InjectNH(ni string, nhIdx uint64) error {
+	niR, ok := f.r.NetworkInstanceRIB(ni)
+	if !ok {
+		return fmt.Errorf("unknown NI, %s", ni)
+	}
+
+	if _, _, err := niR.AddNextHop(&aftpb.Afts_NextHopKey{
+		Index:   nhIdx,
+		NextHop: &aftpb.Afts_NextHop{},
+	}, false); err != nil {
+		return fmt.Errorf("cannot add NH entry, err: %v", err)
+	}
+
+	return nil
+}

rib/helpers_test.go

@@ -27,7 +27,7 @@ import (
 	"github.com/openconfig/ygot/ygot"
 )
 
-func TestRIBFromGetResponses(t *testing.T) {
+func TestFromGetResponses(t *testing.T) {
 	defaultName := "DEFAULT"
 	tests := []struct {
 		desc          string
@@ -240,9 +240,9 @@ func TestRIBFromGetResponses(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.desc, func(t *testing.T) {
-			got, err := RIBFromGetResponses(tt.inDefaultName, tt.inResponses)
+			got, err := FromGetResponses(tt.inDefaultName, tt.inResponses)
 			if (err != nil) != tt.wantErr {
-				t.Fatalf("RIBFromGetResponses(...): did not get expected error, got: %v, wantErr? %v", err, tt.wantErr)
+				t.Fatalf("FromGetResponses(...): did not get expected error, got: %v, wantErr? %v", err, tt.wantErr)
 			}
 
 			if diff := cmp.Diff(got, tt.wantRIB,
@@ -251,7 +251,154 @@ func TestRIBFromGetResponses(t *testing.T) {
 				cmp.AllowUnexported(RIBHolder{}),
 				cmpopts.IgnoreFields(RIBHolder{}, "mu", "refCounts", "checkFn"),
 			); diff != "" {
-				t.Fatalf("RIBFromGetResponses(...): did not get expected RIB, diff(-got,+want):\n%s", diff)
+				t.Fatalf("FromGetResponses(...): did not get expected RIB, diff(-got,+want):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestFakeRIB(t *testing.T) {
+	dn := "DEFAULT"
+	tests := []struct {
+		desc    string
+		inBuild func() *fakeRIB
+		wantRIB *RIB
+	}{{
+		desc: "nh only",
+		inBuild: func() *fakeRIB {
+			f := NewFake(dn)
+			if err := f.InjectNH(dn, 1); err != nil {
+				t.Fatalf("cannot add NH, err: %v", err)
+			}
+			return f
+		},
+		wantRIB: &RIB{
+			defaultName: dn,
+			niRIB: map[string]*RIBHolder{
+				dn: {
+					name: dn,
+					r: &aft.RIB{
+						Afts: &aft.Afts{
+							NextHop: map[uint64]*aft.Afts_NextHop{
+								1: {Index: ygot.Uint64(1)},
+							},
+						},
+					},
+				},
+			},
+		},
+	}, {
+		desc: "ipv4 only",
+		inBuild: func() *fakeRIB {
+			f := NewFake(dn, DisableRIBCheckFn())
+			if err := f.InjectIPv4(dn, "1.0.0.0/24", 1); err != nil {
+				t.Fatalf("cannot add IPv4, err: %v", err)
+			}
+			return f
+		},
+		wantRIB: &RIB{
+			defaultName: dn,
+			niRIB: map[string]*RIBHolder{
+				dn: {
+					name: dn,
+					r: &aft.RIB{
+						Afts: &aft.Afts{
+							Ipv4Entry: map[string]*aft.Afts_Ipv4Entry{
+								"1.0.0.0/24": {Prefix: ygot.String("1.0.0.0/24"), NextHopGroup: ygot.Uint64(1)},
+							},
+						},
+					},
+				},
+			},
+		},
+	}, {
+		desc: "nhg only",
+		inBuild: func() *fakeRIB {
+			f := NewFake(dn, DisableRIBCheckFn())
+			if err := f.InjectNHG(dn, 1, map[uint64]uint64{1: 1}); err != nil {
+				t.Fatalf("cannot add NHG, err: %v", err)
+			}
+			return f
+		},
+		wantRIB: &RIB{
+			defaultName: dn,
+			niRIB: map[string]*RIBHolder{
+				dn: {
+					name: dn,
+					r: &aft.RIB{
+						Afts: &aft.Afts{
+							NextHopGroup: map[uint64]*aft.Afts_NextHopGroup{
+								1: {
+									Id: ygot.Uint64(1),
+									NextHop: map[uint64]*aft.Afts_NextHopGroup_NextHop{
+										1: {
+											Index:  ygot.Uint64(1),
+											Weight: ygot.Uint64(1),
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}, {
+		desc: "full ipv4 chain",
+		inBuild: func() *fakeRIB {
+			f := NewFake(dn)
+			// Discard the errors, since the test will check whether the entries are there.
+			f.InjectNH(dn, 1)
+			f.InjectNHG(dn, 1, map[uint64]uint64{1: 1})
+			f.InjectIPv4(dn, "192.0.2.1/32", 1)
+			return f
+		},
+		wantRIB: &RIB{
+			defaultName: dn,
+			niRIB: map[string]*RIBHolder{
+				dn: {
+					name: dn,
+					r: &aft.RIB{
+						Afts: &aft.Afts{
+							NextHop: map[uint64]*aft.Afts_NextHop{
+								1: {
+									Index: ygot.Uint64(1),
+								},
+							},
+							NextHopGroup: map[uint64]*aft.Afts_NextHopGroup{
+								1: {
+									Id: ygot.Uint64(1),
+									NextHop: map[uint64]*aft.Afts_NextHopGroup_NextHop{
+										1: {
+											Index:  ygot.Uint64(1),
+											Weight: ygot.Uint64(1),
+										},
+									},
+								},
+							},
+							Ipv4Entry: map[string]*aft.Afts_Ipv4Entry{
+								"192.0.2.1/32": {
+									Prefix:       ygot.String("192.0.2.1/32"),
+									NextHopGroup: ygot.Uint64(1),
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			got := tt.inBuild().RIB()
+			if diff := cmp.Diff(got, tt.wantRIB,
+				cmpopts.EquateEmpty(), cmp.AllowUnexported(RIB{}),
+				cmpopts.IgnoreFields(RIB{}, "nrMu", "pendMu", "ribCheck"),
+				cmp.AllowUnexported(RIBHolder{}),
+				cmpopts.IgnoreFields(RIBHolder{}, "mu", "refCounts", "checkFn"),
+			); diff != "" {
+				t.Fatalf("FakeRIB.RIB(...): did not get expected RIB, diff(-got,+want):\n%s", diff)
 			}
 		})
 	}