1 files changed, 18 insertions, 9 deletions
diff --git a/swarm/storage/filestore_test.go b/swarm/storage/filestore_test.go
index 1aaec5e5c..f3f597255 100644
--- a/swarm/storage/filestore_test.go
+++ b/swarm/storage/filestore_test.go
@@ -18,6 +18,7 @@ package storage
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"io/ioutil"
 	"os"
@@ -49,12 +50,16 @@ func testFileStoreRandom(toEncrypt bool, t *testing.T) {
 	defer os.RemoveAll("/tmp/bzz")
 
 	reader, slice := generateRandomData(testDataSize)
-	key, wait, err := fileStore.Store(reader, testDataSize, toEncrypt)
+	ctx := context.TODO()
+	key, wait, err := fileStore.Store(ctx, reader, testDataSize, toEncrypt)
 	if err != nil {
 		t.Errorf("Store error: %v", err)
 	}
-	wait()
-	resultReader, isEncrypted := fileStore.Retrieve(key)
+	err = wait(ctx)
+	if err != nil {
+		t.Fatalf("Store waitt error: %v", err.Error())
+	}
+	resultReader, isEncrypted := fileStore.Retrieve(context.TODO(), key)
 	if isEncrypted != toEncrypt {
 		t.Fatalf("isEncrypted expected %v got %v", toEncrypt, isEncrypted)
 	}
@@ -72,7 +77,7 @@ func testFileStoreRandom(toEncrypt bool, t *testing.T) {
 	ioutil.WriteFile("/tmp/slice.bzz.16M", slice, 0666)
 	ioutil.WriteFile("/tmp/result.bzz.16M", resultSlice, 0666)
 	localStore.memStore = NewMemStore(NewDefaultStoreParams(), db)
-	resultReader, isEncrypted = fileStore.Retrieve(key)
+	resultReader, isEncrypted = fileStore.Retrieve(context.TODO(), key)
 	if isEncrypted != toEncrypt {
 		t.Fatalf("isEncrypted expected %v got %v", toEncrypt, isEncrypted)
 	}
@@ -110,12 +115,16 @@ func testFileStoreCapacity(toEncrypt bool, t *testing.T) {
 	}
 	fileStore := NewFileStore(localStore, NewFileStoreParams())
 	reader, slice := generateRandomData(testDataSize)
-	key, wait, err := fileStore.Store(reader, testDataSize, toEncrypt)
+	ctx := context.TODO()
+	key, wait, err := fileStore.Store(ctx, reader, testDataSize, toEncrypt)
+	if err != nil {
+		t.Errorf("Store error: %v", err)
+	}
+	err = wait(ctx)
 	if err != nil {
 		t.Errorf("Store error: %v", err)
 	}
-	wait()
-	resultReader, isEncrypted := fileStore.Retrieve(key)
+	resultReader, isEncrypted := fileStore.Retrieve(context.TODO(), key)
 	if isEncrypted != toEncrypt {
 		t.Fatalf("isEncrypted expected %v got %v", toEncrypt, isEncrypted)
 	}
@@ -134,7 +143,7 @@ func testFileStoreCapacity(toEncrypt bool, t *testing.T) {
 	memStore.setCapacity(0)
 	// check whether it is, indeed, empty
 	fileStore.ChunkStore = memStore
-	resultReader, isEncrypted = fileStore.Retrieve(key)
+	resultReader, isEncrypted = fileStore.Retrieve(context.TODO(), key)
 	if isEncrypted != toEncrypt {
 		t.Fatalf("isEncrypted expected %v got %v", toEncrypt, isEncrypted)
 	}
@@ -144,7 +153,7 @@ func testFileStoreCapacity(toEncrypt bool, t *testing.T) {
 	// check how it works with localStore
 	fileStore.ChunkStore = localStore
 	//	localStore.dbStore.setCapacity(0)
-	resultReader, isEncrypted = fileStore.Retrieve(key)
+	resultReader, isEncrypted = fileStore.Retrieve(context.TODO(), key)
 	if isEncrypted != toEncrypt {
 		t.Fatalf("isEncrypted expected %v got %v", toEncrypt, isEncrypted)
 	}