better batch testing stuff

flatfs/flatfs.go

@@ -330,12 +330,12 @@ type flatfsBatch struct {
 	ds *Datastore
 }
 
-func (fs *Datastore) Batch() datastore.Batch {
+func (fs *Datastore) Batch() (datastore.Batch, error) {
 	return &flatfsBatch{
 		puts:    make(map[datastore.Key]interface{}),
 		deletes: make(map[datastore.Key]struct{}),
 		ds:      fs,
-	}
+	}, nil
 }
 
 func (bt *flatfsBatch) Put(key datastore.Key, val interface{}) error {

flatfs/flatfs_test.go

 package flatfs_test
 
 import (
-	"bytes"
 	"encoding/base32"
 	"io/ioutil"
 	"os"
@@ -12,6 +11,7 @@ import (
 	"github.com/jbenet/go-datastore"
 	"github.com/jbenet/go-datastore/flatfs"
 	"github.com/jbenet/go-datastore/query"
+	dstest "github.com/jbenet/go-datastore/test"
 
 	rand "github.com/jbenet/go-datastore/Godeps/_workspace/src/github.com/dustin/randbo"
 )
@@ -330,38 +330,19 @@ func TestBatchPut(t *testing.T) {
 		t.Fatalf("New fail: %v\n", err)
 	}
 
-	batch := fs.Batch()
-	r := rand.New()
-	var blocks [][]byte
-	var keys []datastore.Key
-	for i := 0; i < 20; i++ {
-		blk := make([]byte, 256*1024)
-		r.Read(blk)
-		blocks = append(blocks, blk)
+	dstest.RunBatchTest(t, fs)
+}
 
-		key := datastore.NewKey(base32.StdEncoding.EncodeToString(blk[:8]))
-		keys = append(keys, key)
+func TestBatchDelete(t *testing.T) {
+	temp, cleanup := tempdir(t)
+	defer cleanup()
 
-		err := batch.Put(key, blk)
-		if err != nil {
-			t.Fatal(err)
-		}
-	}
-	err = batch.Commit()
+	fs, err := flatfs.New(temp, 2)
 	if err != nil {
-		t.Fatal(err)
+		t.Fatalf("New fail: %v\n", err)
 	}
 
-	for i, k := range keys {
-		blk, err := fs.Get(k)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		if !bytes.Equal(blk.([]byte), blocks[i]) {
-			t.Fatal("blocks not correct!")
-		}
-	}
+	dstest.RunBatchDeleteTest(t, fs)
 }
 
 func BenchmarkConsecutivePut(b *testing.B) {
@@ -417,7 +398,10 @@ func BenchmarkBatchedPut(b *testing.B) {
 	b.ResetTimer()
 
 	for i := 0; i < b.N; {
-		batch := fs.Batch()
+		batch, err := fs.Batch()
+		if err != nil {
+			b.Fatal(err)
+		}
 
 		for n := i; i-n < 512 && i < b.N; i++ {
 			err := batch.Put(keys[i], blocks[i])

batch_test.go → test/test_util.go

-package datastore
+package dstest
 
 import (
 	"bytes"
 	"encoding/base32"
 	"testing"
 
+	dstore "github.com/jbenet/go-datastore"
 	rand "github.com/jbenet/go-datastore/Godeps/_workspace/src/github.com/dustin/randbo"
 )
 
-func TestBasicBatch(t *testing.T) {
-	ds := NewMapDatastore()
-	batch := NewBasicBatch(ds)
+func RunBatchTest(t *testing.T, ds dstore.BatchingDatastore) {
+	batch, err := ds.Batch()
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	r := rand.New()
 	var blocks [][]byte
-	var keys []Key
+	var keys []dstore.Key
 	for i := 0; i < 20; i++ {
 		blk := make([]byte, 256*1024)
 		r.Read(blk)
 		blocks = append(blocks, blk)
 
-		key := NewKey(base32.StdEncoding.EncodeToString(blk[:8]))
+		key := dstore.NewKey(base32.StdEncoding.EncodeToString(blk[:8]))
 		keys = append(keys, key)
 
 		err := batch.Put(key, blk)
@@ -28,7 +31,17 @@ func TestBasicBatch(t *testing.T) {
 			t.Fatal(err)
 		}
 	}
-	err := batch.Commit()
+
+	// Ensure they are not in the datastore before comitting
+	for _, k := range keys {
+		_, err := ds.Get(k)
+		if err == nil {
+			t.Fatal("should not have found this block")
+		}
+	}
+
+	// commit, write them to the datastore
+	err = batch.Commit()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -45,16 +58,14 @@ func TestBasicBatch(t *testing.T) {
 	}
 }
 
-func TestBasicBatchDelete(t *testing.T) {
-	ds := NewMapDatastore()
-
+func RunBatchDeleteTest(t *testing.T, ds dstore.BatchingDatastore) {
 	r := rand.New()
-	var keys []Key
+	var keys []dstore.Key
 	for i := 0; i < 20; i++ {
 		blk := make([]byte, 16)
 		r.Read(blk)
 
-		key := NewKey(base32.StdEncoding.EncodeToString(blk[:8]))
+		key := dstore.NewKey(base32.StdEncoding.EncodeToString(blk[:8]))
 		keys = append(keys, key)
 
 		err := ds.Put(key, blk)
@@ -63,14 +74,18 @@ func TestBasicBatchDelete(t *testing.T) {
 		}
 	}
 
-	batch := NewBasicBatch(ds)
+	batch, err := ds.Batch()
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	for _, k := range keys {
 		err := batch.Delete(k)
 		if err != nil {
 			t.Fatal(err)
 		}
 	}
-	err := batch.Commit()
+	err = batch.Commit()
 	if err != nil {
 		t.Fatal(err)
 	}