logdog: Add bundler library.

client/internal/logdog/butler/bundler/bundler_test.go

+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package bundler
+
+import (
+        "crypto/md5"
+        "encoding/hex"
+        "fmt"
+        "strings"
+        "testing"
+        "time"
+
+        "github.com/golang/protobuf/proto"
+        "github.com/luci/luci-go/common/logdog/protocol"
+        "github.com/luci/luci-go/common/logdog/protocol/protoutil"
+        "github.com/luci/luci-go/common/logdog/types"
+        . "github.com/smartystreets/goconvey/convey"
+)
+
+var (
+        testNow = time.Date(2015, 1, 1, 0, 0, 0, 0, time.UTC)
+)
+
+// fakeSizer is a Sizer implementation that counts (obviously incorrect) fixed
+// sizes for each entry type.
+type fakeSizer struct {
+        Bundle      int64
+        BundleEntry int64
+        LogEntry    int64
+
+        size int64
+        seen map[types.StreamPath]bool
+
+        lastSize int64
+        lastSeen types.StreamPath
+}
+
+func (s *fakeSizer) Size() int64 {
+        return s.Bundle + s.size
+}
+
+func (s *fakeSizer) Append(be *protocol.ButlerLogBundle_Entry, e *protocol.LogEntry) {
+        size := int64(0)
+
+        // Add the ButlerLogBundle_Entry code if we haven't seen it before.
+        path := protoutil.DescriptorPath(be.GetDesc())
+        if seen := s.seen[path]; !seen {
+                if s.seen == nil {
+                        s.seen = map[types.StreamPath]bool{
+                                path: true,
+                        }
+                } else {
+                        s.seen[path] = true
+                }
+                s.lastSeen = path
+                size += s.BundleEntry
+        } else {
+                s.lastSeen = ""
+        }
+
+        if e != nil {
+                // Each character in the line gets LogEntry space.
+                if len(e.GetLines()) == 1 {
+                        size += s.LogEntry * int64(len(e.GetLines()[0]))
+                } else {
+                        size += s.LogEntry
+                }
+        }
+
+        s.lastSize = size
+        s.size += size
+}
+
+func (s *fakeSizer) Undo() {
+        s.size -= s.lastSize
+        if s.lastSeen != "" {
+                s.seen[s.lastSeen] = false
+        }
+}
+
+func hash(s, t string) []byte {
+        sum := md5.Sum([]byte(fmt.Sprintf("%s::%s", s, t)))
+        return sum[:]
+}
+
+func key(s, t string) string {
+        return hex.EncodeToString(hash(s, t))
+}
+
+// addEntry generates a ButlerLogBundle_Entry and appends it to our Bundler via
+// one or more calls to Append.
+//
+// If "le" strings are supplied, those will create generated LogEntry for that
+// ButlerLogBundle_Entry.
+func gen(e string, t bool, le ...string) *protocol.ButlerLogBundle_Entry {
+        secret := hash(e, "secret")
+        name := key(e, "name")
+        contentType := "test/data"
+
+        be := &protocol.ButlerLogBundle_Entry{
+                Desc: &protocol.LogStreamDescriptor{
+                        Prefix:      &e,
+                        Name:        &name,
+                        ContentType: &contentType,
+                        Timestamp:   protoutil.NewTimestamp(testNow),
+                },
+                Secret:   secret,
+                Terminal: &t,
+        }
+
+        if len(le) > 0 {
+                be.Logs = make([]*protocol.LogEntry, len(le))
+                for i, l := range le {
+                        be.Logs[i] = &protocol.LogEntry{
+                                Lines: []string{
+                                        l,
+                                },
+                                Data: [][]byte{
+                                        hash(l, "data0"),
+                                        hash(l, "data1"),
+                                        hash(l, "data2"),
+                                },
+                        }
+                }
+        }
+        return be
+}
+
+func logEntryName(le *protocol.LogEntry) string {
+        if len(le.GetLines()) != 1 {
+                return ""
+        }
+        return le.GetLines()[0]
+}
+
+// "expected" is a notation to express a bundle entry and its keys:
+//   "a": a bundle entry keyed on "a".
+//   "+a": a terminal bundle entry keyed on "a".
+//   "a:1:2:3": a bundle entry keyed on "a" with three log entries, each keyed on
+//            "1", "2", and "3" respectively.
+func shouldHaveBundleEntries(actual interface{}, expected ...interface{}) string {
+        bundle := actual.(*protocol.ButlerLogBundle)
+
+        errors := []string{}
+        fail := func(f string, args ...interface{}) {
+                errors = append(errors, fmt.Sprintf(f, args...))
+        }
+
+        term := make(map[string]bool)
+        exp := make(map[string][]string)
+
+        // Parse expectation strings.
+        for _, e := range expected {
+                s := e.(string)
+                if len(s) == 0 {
+                        continue
+                }
+
+                t := false
+                if s[0] == '+' {
+                        t = true
+                        s = s[1:]
+                }
+
+                parts := strings.Split(s, ":")
+                name := parts[0]
+                term[name] = t
+
+                if len(parts) > 1 {
+                        exp[name] = append(exp[name], parts[1:]...)
+                }
+        }
+
+        entries := make(map[string]*protocol.ButlerLogBundle_Entry)
+        for _, be := range bundle.GetEntries() {
+                entries[be.GetDesc().GetPrefix()] = be
+        }
+        for name, t := range term {
+                be := entries[name]
+                if be == nil {
+                        fail("No bundle entry for [%s]", name)
+                        continue
+                }
+                delete(entries, name)
+
+                if t != be.GetTerminal() {
+                        fail("Bundle entry [%s] doesn't match expected terminal state (exp: %v != act: %v)",
+                                name, t, be.GetTerminal())
+                }
+
+                logs := exp[name]
+                for i, l := range logs {
+                        if i >= len(be.GetLogs()) {
+                                fail("Bundle entry [%s] missing log: %s", name, l)
+                                continue
+                        }
+                        le := be.GetLogs()[i]
+
+                        if logEntryName(le) != l {
+                                fail("Bundle entry [%s] log %d doesn't match expected (exp: %s != act: %s)",
+                                        name, i, l, logEntryName(le))
+                                continue
+                        }
+                }
+                if len(be.GetLogs()) > len(logs) {
+                        for _, le := range be.GetLogs()[len(logs):] {
+                                fail("Bundle entry [%s] has extra log entry: %s", name, logEntryName(le))
+                        }
+                }
+        }
+        for k := range entries {
+                fail("Unexpected bundle entry present: [%s]", k)
+        }
+        return strings.Join(errors, "\n")
+}
+
+func TestBundler(t *testing.T) {
+        Convey(`An empty Bundler`, t, func() {
+                b := New(Config{}).(*bundlerImpl)
+
+                Convey(`Has a size of 0 and nil GetBundles() return value.`, func() {
+                        So(b.Size(), ShouldEqual, 0)
+                        So(b.Empty(), ShouldBeTrue)
+                        So(b.GetBundles(), ShouldBeNil)
+                })
+
+                Convey(`When adding an empty entry, still has size 0 and nil GetBundles() return value.`, func() {
+                        b.Append(gen("a", false))
+                        So(b.Size(), ShouldEqual, 0)
+                        So(b.Empty(), ShouldBeTrue)
+                        So(b.GetBundles(), ShouldBeNil)
+                })
+
+                Convey(`Bundles a terminal entry with no logs.`, func() {
+                        b.Append(gen("a", true))
+
+                        size, empty, bundles := b.Size(), b.Empty(), b.GetBundles()
+                        So(empty, ShouldBeFalse)
+
+                        So(len(bundles), ShouldEqual, 1)
+                        So(size, ShouldBeGreaterThanOrEqualTo, proto.Size(bundles[0]))
+                        So(bundles[0], shouldHaveBundleEntries, "+a")
+                })
+
+                Convey(`Bundles an entry with 3 logs.`, func() {
+                        b.Append(gen("a", false, "1", "2"))
+                        b.Append(gen("a", false, "3"))
+
+                        size, empty, bundles := b.Size(), b.Empty(), b.GetBundles()
+                        So(empty, ShouldBeFalse)
+
+                        So(len(bundles), ShouldEqual, 1)
+                        So(size, ShouldBeGreaterThanOrEqualTo, proto.Size(bundles[0]))
+                        So(bundles[0], shouldHaveBundleEntries, "a:1:2:3")
+                })
+
+                Convey(`Bundles 2 entries with 2 logs each and one terminal entry with no logs.`, func() {
+                        b.Append(gen("a", false, "1", "2"))
+                        b.Append(gen("b", false, "3", "4"))
+                        b.Append(gen("c", true))
+                        b.Append(gen("d", false))
+
+                        size, empty, bundles := b.Size(), b.Empty(), b.GetBundles()
+                        So(empty, ShouldBeFalse)
+
+                        So(len(bundles), ShouldEqual, 1)
+                        So(size, ShouldBeGreaterThanOrEqualTo, proto.Size(bundles[0]))
+                        So(bundles[0], shouldHaveBundleEntries, "a:1:2", "b:3:4", "+c")
+                })
+        })
+
+        Convey(`A Bundler with a fake Sizer`, t, func() {
+                source := "test suite"
+                b := New(Config{
+                        Threshold: 20,
+                        TemplateBundle: protocol.ButlerLogBundle{
+                                Source: &source,
+                        },
+                        NewSizer: func(*protocol.ButlerLogBundle) Sizer {
+                                return &fakeSizer{
+                                        Bundle:      8,
+                                        BundleEntry: 2,
+                                        LogEntry:    5,
+                                }
+                        },
+                }).(*bundlerImpl)
+                So(b.Size(), ShouldEqual, 8)
+
+                Convey(`Adding an entry with 5 log messages outputs three bundles.`, func() {
+                        b.Append(gen("a", false, "1", "2", "3", "4", "5"))
+
+                        bundles := b.GetBundles()
+                        So(len(bundles), ShouldEqual, 3)
+
+                        Convey(`All bundles use the template bundle's fields.`, func() {
+                                So(bundles[0].GetSource(), ShouldEqual, source)
+                                So(bundles[1].GetSource(), ShouldEqual, source)
+                                So(bundles[2].GetSource(), ShouldEqual, source)
+                        })
+
+                        Convey(`Have the right entries: {1,2}, {3,4}, {5}.`, func() {
+                                So(bundles[0], shouldHaveBundleEntries, "a:1:2")
+                                So(bundles[1], shouldHaveBundleEntries, "a:3:4")
+                                So(bundles[2], shouldHaveBundleEntries, "a:5")
+                        })
+                })
+
+                Convey(`Adding two entries with 2 log messages each outputs first, then second.`, func() {
+                        b.Append(gen("a", false, "1", "2"))
+                        b.Append(gen("b", false, "3", "4"))
+
+                        bundles := b.GetBundles()
+                        So(len(bundles), ShouldEqual, 2)
+                        So(bundles[0], shouldHaveBundleEntries, "a:1:2")
+                        So(bundles[1], shouldHaveBundleEntries, "b:3:4")
+                })
+
+                Convey(`A non-terminal entry followed by a terminal version gets output as terminal.`, func() {
+                        b.Append(gen("a", false, "1"))
+                        b.Append(gen("a", true, "2"))
+
+                        bundles := b.GetBundles()
+                        So(len(bundles), ShouldEqual, 1)
+                        So(bundles[0], shouldHaveBundleEntries, "+a:1:2")
+                })
+
+                Convey(`A terminal entry followed by a non-terminal version gets output as terminal.`, func() {
+                        b.Append(gen("a", true, "1"))
+                        b.Append(gen("a", false, "2"))
+                        b.Append(gen("a", false))
+
+                        bundles := b.GetBundles()
+                        So(len(bundles), ShouldEqual, 1)
+                        So(bundles[0], shouldHaveBundleEntries, "+a:1:2")
+                })
+
+                Convey(`When the base bundle is above threshold, clears logs and returns nil.`, func() {
+                        b.Append(gen("a", true))
+
+                        So(b.Size(), ShouldEqual, 10)
+                        So(b.getBundlesImpl(7), ShouldBeNil)
+
+                        So(b.Size(), ShouldEqual, 8)
+                        So(b.getBundlesImpl(0), ShouldBeNil)
+                })
+
+                Convey(`When the bundle entry size is above threshold, clears logs and returns nil.`, func() {
+                        b.Append(gen("a", true))
+
+                        So(b.Size(), ShouldEqual, 10)
+                        So(b.getBundlesImpl(9), ShouldBeNil)
+
+                        So(b.Size(), ShouldEqual, 8)
+                        So(b.getBundlesImpl(0), ShouldBeNil)
+                })
+
+                Convey(`When the bundle has a log entry larger than threshold, it discards it.`, func() {
+                        b.Append(gen("a", false, "a", "bbb", "cc", "d"))
+                        b.Append(gen("b", false, "1", "2", "3"))
+
+                        bundles := b.GetBundles()
+                        So(len(bundles), ShouldEqual, 4)
+                        So(bundles[0], shouldHaveBundleEntries, "a:a:d")
+                        So(bundles[1], shouldHaveBundleEntries, "a:cc")
+                        So(bundles[2], shouldHaveBundleEntries, "b:1:2")
+                        So(bundles[3], shouldHaveBundleEntries, "b:3")
+
+                        So(b.getBundlesImpl(0), ShouldBeNil)
+                })
+        })
+}