From 995c4ec88093e1e15bcb8689f4c718f064156c67 Mon Sep 17 00:00:00 2001
From: Kevin Atkinson <k@kevina.org>
Date: Sun, 12 Feb 2017 23:02:29 -0500
Subject: [PATCH] Add test for EnumerateChildren progress reporting.

License: MIT
Signed-off-by: Kevin Atkinson <k@kevina.org>
---
 merkledag/merkledag_test.go | 79 +++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/merkledag/merkledag_test.go b/merkledag/merkledag_test.go
index c55a7b551f3..d7c7df77df8 100644
--- a/merkledag/merkledag_test.go
+++ b/merkledag/merkledag_test.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"math/rand"
 	"strings"
 	"sync"
 	"testing"
@@ -547,3 +548,81 @@ func TestEnumerateAsyncFailsNotFound(t *testing.T) {
 		t.Fatal("this should have failed")
 	}
 }
+
+func TestProgressIndicator(t *testing.T) {
+	testProgressIndicator(t, 5)
+}
+
+func TestProgressIndicatorNoChildren(t *testing.T) {
+	testProgressIndicator(t, 0)
+}
+
+func testProgressIndicator(t *testing.T, depth int) {
+	ds := dstest.Mock()
+
+	top, numChildren := mkDag(ds, depth)
+	cset := cid.NewSet()
+
+	v := new(ProgressTracker)
+	ctx := context.WithValue(context.Background(), "progress", v)
+
+	err := EnumerateChildren(ctx, ds, top, cset.Visit, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if v.Total != numChildren {
+		t.Errorf("wrong number of children reported in progress indicator, expected %d, got %d",
+			numChildren, v.Total)
+	}
+}
+
+func mkDag(ds DAGService, depth int) (*cid.Cid, int) {
+	totalChildren := 0
+	f := func() *ProtoNode {
+		p := new(ProtoNode)
+		buf := make([]byte, 16)
+		rand.Read(buf)
+
+		p.SetData(buf)
+		_, err := ds.Add(p)
+		if err != nil {
+			panic(err)
+		}
+		return p
+	}
+
+	for i := 0; i < depth; i++ {
+		thisf := f
+		f = func() *ProtoNode {
+			pn := mkNodeWithChildren(thisf, 10)
+			_, err := ds.Add(pn)
+			if err != nil {
+				panic(err)
+			}
+			totalChildren += 10
+			return pn
+		}
+	}
+
+	nd := f()
+	c, err := ds.Add(nd)
+	if err != nil {
+		panic(err)
+	}
+
+	return c, totalChildren
+}
+
+func mkNodeWithChildren(getChild func() *ProtoNode, width int) *ProtoNode {
+	cur := new(ProtoNode)
+
+	for i := 0; i < width; i++ {
+		c := getChild()
+		if err := cur.AddNodeLinkClean(fmt.Sprint(i), c); err != nil {
+			panic(err)
+		}
+	}
+
+	return cur
+}