zipkey: add batch test

1 files changed, 57 insertions, 0 deletions

diff --git a/skate/zipkey/batch_test.go b/skate/zipkey/batch_test.go
index 7c6a48c..38a1307 100644
--- a/skate/zipkey/batch_test.go
+++ b/skate/zipkey/batch_test.go
@@ -3,7 +3,11 @@ package zipkey
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
+	"io"
+	"reflect"
 	"strings"
+	"sync"
 	"testing"
 )
 
@@ -30,3 +34,56 @@ func TestBatcher(t *testing.T) {
 		t.Fatalf("got %v, want %v", got, want)
 	}
 }
+
+func TestBatcherLarge(t *testing.T) {
+	var (
+		N          = 1000000
+		numWorkers = 24
+		size       = 7000
+		// We share a single writer across threads, so we need to guard each
+		// write. TODO: measure performance impact.
+		mu  sync.Mutex
+		buf bytes.Buffer
+		f   = func(g *Group) error {
+			var v string
+			if reflect.DeepEqual(g.G0, g.G1) {
+				v = "1"
+			} else {
+				v = "0"
+			}
+			mu.Lock()
+			defer mu.Unlock()
+			if _, err := io.WriteString(&buf, v); err != nil {
+				return err
+			}
+			return nil
+		}
+		b = NewBatcher(groupFunc(f))
+	)
+	b.Size = size
+	b.NumWorkers = numWorkers
+	for i := 0; i < N; i++ {
+		var u, v string
+		if i%2 == 0 {
+			u, v = "a", "b"
+		} else {
+			u, v = "a", "a"
+		}
+		g := &Group{
+			Key: fmt.Sprintf("%d", i),
+			G0:  []string{u},
+			G1:  []string{v},
+		}
+		if err := b.GroupFunc(g); err != nil {
+			t.Fatalf("unexpected err from gf: %v", err)
+		}
+	}
+	if err := b.Close(); err != nil {
+		t.Fatalf("unexpected err from close: %v", err)
+	}
+	got := buf.String()
+	count0, count1 := strings.Count(got, "0"), strings.Count(got, "1")
+	if count1 != N/2 {
+		t.Fatalf("got %v, want %v (count0=%v, buf=%s)", count1, N/2, count0, buf.String())
+	}
+}