diff --git a/internal/transport/gnmiext/client.go b/internal/transport/gnmiext/client.go index 037b92fe1..f05fc9eea 100644 --- a/internal/transport/gnmiext/client.go +++ b/internal/transport/gnmiext/client.go @@ -116,6 +116,21 @@ func (b *SetBuilder) Limit(n int) *SetBuilder { return b } +// Split partitions operations by a predicate. Matching ops go into the +// first builder, the rest into the second. Both inherit the limit. +func (b *SetBuilder) Split(fn func(DataElement) bool) (match, rest *SetBuilder) { + match = &SetBuilder{limit: b.limit} + rest = &SetBuilder{limit: b.limit} + for _, op := range b.ops { + if fn(op.el) { + match.ops = append(match.ops, op) + } else { + rest.ops = append(rest.ops, op) + } + } + return match, rest +} + //go:generate go tool moq -with-resets -out client_mock.go . Client type Client interface { diff --git a/internal/transport/gnmiext/client_test.go b/internal/transport/gnmiext/client_test.go index fb99af027..244db94b1 100644 --- a/internal/transport/gnmiext/client_test.go +++ b/internal/transport/gnmiext/client_test.go @@ -1273,6 +1273,57 @@ func TestClient_Do(t *testing.T) { } } +func TestSetBuilder_Split(t *testing.T) { + a := new(Hostname("a")) + b := new(Hostname("b")) + c := new(DefaultableHostname("c")) + d := new(Hostname("d")) + + builder := new(SetBuilder).Limit(5). + Update(a). + Patch(b). + Delete(c). + Update(d) + + match, rest := builder.Split(func(el DataElement) bool { + _, ok := el.(*DefaultableHostname) + return ok + }) + + // match should contain only the DefaultableHostname + if len(match.ops) != 1 { + t.Fatalf("Split() match: got %d ops, want 1", len(match.ops)) + } + if match.ops[0].el != c { + t.Errorf("Split() match[0]: got %v, want %v", match.ops[0].el, c) + } + if match.ops[0].mode != del { + t.Errorf("Split() match[0].mode: got %v, want del", match.ops[0].mode) + } + + // rest should contain the other three + if len(rest.ops) != 3 { + t.Fatalf("Split() rest: got %d ops, want 3", len(rest.ops)) + } + if rest.ops[0].el != a || rest.ops[0].mode != replace { + t.Errorf("Split() rest[0]: got %v/%v, want a/replace", rest.ops[0].el, rest.ops[0].mode) + } + if rest.ops[1].el != b || rest.ops[1].mode != update { + t.Errorf("Split() rest[1]: got %v/%v, want b/update", rest.ops[1].el, rest.ops[1].mode) + } + if rest.ops[2].el != d || rest.ops[2].mode != replace { + t.Errorf("Split() rest[2]: got %v/%v, want d/replace", rest.ops[2].el, rest.ops[2].mode) + } + + // Both inherit the limit + if match.limit != 5 { + t.Errorf("Split() match.limit: got %d, want 5", match.limit) + } + if rest.limit != 5 { + t.Errorf("Split() rest.limit: got %d, want 5", rest.limit) + } +} + func TestStringToStructuredPath(t *testing.T) { tests := []struct { name string