1 files changed, 45 insertions, 22 deletions

diff --git a/swarm/storage/chunker_test.go b/swarm/storage/chunker_test.go
index d8be13ef6..69c388b39 100644
--- a/swarm/storage/chunker_test.go
+++ b/swarm/storage/chunker_test.go
@@ -18,6 +18,7 @@ package storage
 
 import (
 	"bytes"
+	"context"
 	"crypto/rand"
 	"encoding/binary"
 	"errors"
@@ -81,7 +82,7 @@ func testRandomBrokenData(n int, tester *chunkerTester) {
 	putGetter := newTestHasherStore(NewMapChunkStore(), SHA3Hash)
 
 	expectedError := fmt.Errorf("Broken reader")
-	addr, _, err := TreeSplit(brokendata, int64(n), putGetter)
+	addr, _, err := TreeSplit(context.TODO(), brokendata, int64(n), putGetter)
 	if err == nil || err.Error() != expectedError.Error() {
 		tester.t.Fatalf("Not receiving the correct error! Expected %v, received %v", expectedError, err)
 	}
@@ -104,20 +105,24 @@ func testRandomData(usePyramid bool, hash string, n int, tester *chunkerTester)
 	putGetter := newTestHasherStore(NewMapChunkStore(), hash)
 
 	var addr Address
-	var wait func()
+	var wait func(context.Context) error
 	var err error
+	ctx := context.TODO()
 	if usePyramid {
-		addr, wait, err = PyramidSplit(data, putGetter, putGetter)
+		addr, wait, err = PyramidSplit(ctx, data, putGetter, putGetter)
 	} else {
-		addr, wait, err = TreeSplit(data, int64(n), putGetter)
+		addr, wait, err = TreeSplit(ctx, data, int64(n), putGetter)
 	}
 	if err != nil {
 		tester.t.Fatalf(err.Error())
 	}
 	tester.t.Logf(" Key = %v\n", addr)
-	wait()
+	err = wait(ctx)
+	if err != nil {
+		tester.t.Fatalf(err.Error())
+	}
 
-	reader := TreeJoin(addr, putGetter, 0)
+	reader := TreeJoin(context.TODO(), addr, putGetter, 0)
 	output := make([]byte, n)
 	r, err := reader.Read(output)
 	if r != n || err != io.EOF {
@@ -200,11 +205,15 @@ func TestDataAppend(t *testing.T) {
 		chunkStore := NewMapChunkStore()
 		putGetter := newTestHasherStore(chunkStore, SHA3Hash)
 
-		addr, wait, err := PyramidSplit(data, putGetter, putGetter)
+		ctx := context.TODO()
+		addr, wait, err := PyramidSplit(ctx, data, putGetter, putGetter)
+		if err != nil {
+			tester.t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			tester.t.Fatalf(err.Error())
 		}
-		wait()
 
 		//create a append data stream
 		appendInput, found := tester.inputs[uint64(m)]
@@ -217,13 +226,16 @@ func TestDataAppend(t *testing.T) {
 		}
 
 		putGetter = newTestHasherStore(chunkStore, SHA3Hash)
-		newAddr, wait, err := PyramidAppend(addr, appendData, putGetter, putGetter)
+		newAddr, wait, err := PyramidAppend(ctx, addr, appendData, putGetter, putGetter)
+		if err != nil {
+			tester.t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			tester.t.Fatalf(err.Error())
 		}
-		wait()
 
-		reader := TreeJoin(newAddr, putGetter, 0)
+		reader := TreeJoin(ctx, newAddr, putGetter, 0)
 		newOutput := make([]byte, n+m)
 		r, err := reader.Read(newOutput)
 		if r != (n + m) {
@@ -282,12 +294,16 @@ func benchmarkSplitJoin(n int, t *testing.B) {
 		data := testDataReader(n)
 
 		putGetter := newTestHasherStore(NewMapChunkStore(), SHA3Hash)
-		key, wait, err := PyramidSplit(data, putGetter, putGetter)
+		ctx := context.TODO()
+		key, wait, err := PyramidSplit(ctx, data, putGetter, putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
-		wait()
-		reader := TreeJoin(key, putGetter, 0)
+		err = wait(ctx)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		reader := TreeJoin(ctx, key, putGetter, 0)
 		benchReadAll(reader)
 	}
 }
@@ -298,7 +314,7 @@ func benchmarkSplitTreeSHA3(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, SHA3Hash)
 
-		_, _, err := TreeSplit(data, int64(n), putGetter)
+		_, _, err := TreeSplit(context.TODO(), data, int64(n), putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -311,7 +327,7 @@ func benchmarkSplitTreeBMT(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, BMTHash)
 
-		_, _, err := TreeSplit(data, int64(n), putGetter)
+		_, _, err := TreeSplit(context.TODO(), data, int64(n), putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -324,7 +340,7 @@ func benchmarkSplitPyramidSHA3(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, SHA3Hash)
 
-		_, _, err := PyramidSplit(data, putGetter, putGetter)
+		_, _, err := PyramidSplit(context.TODO(), data, putGetter, putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -338,7 +354,7 @@ func benchmarkSplitPyramidBMT(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, BMTHash)
 
-		_, _, err := PyramidSplit(data, putGetter, putGetter)
+		_, _, err := PyramidSplit(context.TODO(), data, putGetter, putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -354,18 +370,25 @@ func benchmarkSplitAppendPyramid(n, m int, t *testing.B) {
 		chunkStore := NewMapChunkStore()
 		putGetter := newTestHasherStore(chunkStore, SHA3Hash)
 
-		key, wait, err := PyramidSplit(data, putGetter, putGetter)
+		ctx := context.TODO()
+		key, wait, err := PyramidSplit(ctx, data, putGetter, putGetter)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
-		wait()
 
 		putGetter = newTestHasherStore(chunkStore, SHA3Hash)
-		_, wait, err = PyramidAppend(key, data1, putGetter, putGetter)
+		_, wait, err = PyramidAppend(ctx, key, data1, putGetter, putGetter)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
-		wait()
 	}
 }