test: Add test datastore selection support (sourcenetwork#88)

* Move discard to after error check * Fix if-else ok is always false here * Add support for running tests against multiple ds types
shahzadlone · Jan 19, 2022 · bd01c5e · bd01c5e
1 parent b0e3ec4
commit bd01c5e
Show file tree

Hide file tree

Showing 3 changed files with 117 additions and 48 deletions.
diff --git a/db/query.go b/db/query.go
@@ -26,11 +26,11 @@ func (db *DB) ExecQuery(ctx context.Context, query string) *client.QueryResult {
 	}
 
 	txn, err := db.NewTxn(ctx, false)
-	defer txn.Discard(ctx)
 	if err != nil {
 		res.Errors = []interface{}{err.Error()}
 		return res
 	}
+	defer txn.Discard(ctx)
 
 	results, err := db.queryExecutor.ExecQuery(ctx, db, txn, query)
 	if err != nil {

diff --git a/db/tests/utils.go b/db/tests/utils.go
@@ -12,9 +12,11 @@ package tests_test
 import (
 	"context"
 	"fmt"
+	"os"
 	"testing"
 
 	badger "github.com/dgraph-io/badger/v3"
+	ds "github.com/ipfs/go-datastore"
 	"github.com/sourcenetwork/defradb/client"
 	badgerds "github.com/sourcenetwork/defradb/datastores/badger/v3"
 	"github.com/sourcenetwork/defradb/db"
@@ -34,66 +36,137 @@ type QueryTestCase struct {
 	Results []map[string]interface{}
 }
 
-func NewMemoryDB() (*db.DB, error) {
+type databaseInfo struct {
+	name string
+	db   *db.DB
+}
+
+var badgerInMemory bool
+var mapStore bool
+
+func init() {
+	// We use environment variables instead of flags `go test ./...` throws for all packages that don't have the flag defined
+	_, badgerInMemory = os.LookupEnv("DEFRA_BADGER_MEMORY")
+	_, mapStore = os.LookupEnv("DEFRA_MAP")
+
+	// default is to run against all
+	if !badgerInMemory && !mapStore {
+		badgerInMemory = true
+		mapStore = true
+	}
+}
+
+func newBadgerMemoryDB() (databaseInfo, error) {
 	opts := badgerds.Options{Options: badger.DefaultOptions("").WithInMemory(true)}
 	rootstore, err := badgerds.NewDatastore("", &opts)
 	if err != nil {
-		return nil, err
+		return databaseInfo{}, err
+	}
+
+	db, err := db.NewDB(rootstore, struct{}{})
+	if err != nil {
+		return databaseInfo{}, err
+	}
+
+	return databaseInfo{
+		name: "badger-in-memory",
+		db:   db,
+	}, nil
+}
+
+func newMapDB() (databaseInfo, error) {
+	rootstore := ds.NewMapDatastore()
+	db, err := db.NewDB(rootstore, struct{}{})
+	if err != nil {
+		return databaseInfo{}, err
 	}
 
-	return db.NewDB(rootstore, struct{}{})
+	return databaseInfo{
+		name: "ipfs-map-datastore",
+		db:   db,
+	}, nil
+}
+
+func getDatabases() ([]databaseInfo, error) {
+	databases := []databaseInfo{}
+
+	if badgerInMemory {
+		badgerIMDatabase, err := newBadgerMemoryDB()
+		if err != nil {
+			return nil, err
+		}
+		databases = append(databases, badgerIMDatabase)
+	}
+
+	if mapStore {
+		mapDatabase, err := newMapDB()
+		if err != nil {
+			return nil, err
+		}
+		databases = append(databases, mapDatabase)
+	}
+
+	return databases, nil
 }
 
 func ExecuteQueryTestCase(t *testing.T, schema string, collectionNames []string, test QueryTestCase) {
 	ctx := context.Background()
-	db, err := NewMemoryDB()
+	dbs, err := getDatabases()
 	assert.NoError(t, err)
+	assert.NotEmpty(t, dbs)
 
-	err = db.AddSchema(ctx, schema)
-	assert.NoError(t, err)
+	for _, dbi := range dbs {
+		fmt.Println("--------------")
+		//nolint:gosimple
+		fmt.Println(fmt.Sprintf("Running tests with database type: %s", dbi.name))
 
-	collections := []client.Collection{}
-	for _, collectionName := range collectionNames {
-		col, err := db.GetCollection(ctx, collectionName)
+		db := dbi.db
+		err = db.AddSchema(ctx, schema)
 		assert.NoError(t, err)
-		collections = append(collections, col)
-	}
 
-	// insert docs
-	for cid, docs := range test.Docs {
-		for i, docStr := range docs {
-			doc, err := document.NewFromJSON([]byte(docStr))
-			assert.NoError(t, err, test.Description)
-			err = collections[cid].Save(ctx, doc)
-			assert.NoError(t, err, test.Description)
-
-			// check for updates
-			updates, ok := test.Updates[i]
-			if ok {
-				for _, u := range updates {
-					err = doc.SetWithJSON([]byte(u))
-					assert.NoError(t, err, test.Description)
-					err = collections[cid].Save(ctx, doc)
-					assert.NoError(t, err, test.Description)
+		collections := []client.Collection{}
+		for _, collectionName := range collectionNames {
+			col, err := db.GetCollection(ctx, collectionName)
+			assert.NoError(t, err)
+			collections = append(collections, col)
+		}
+
+		// insert docs
+		for cid, docs := range test.Docs {
+			for i, docStr := range docs {
+				doc, err := document.NewFromJSON([]byte(docStr))
+				assert.NoError(t, err, test.Description)
+				err = collections[cid].Save(ctx, doc)
+				assert.NoError(t, err, test.Description)
+
+				// check for updates
+				updates, ok := test.Updates[i]
+				if ok {
+					for _, u := range updates {
+						err = doc.SetWithJSON([]byte(u))
+						assert.NoError(t, err, test.Description)
+						err = collections[cid].Save(ctx, doc)
+						assert.NoError(t, err, test.Description)
+					}
 				}
 			}
 		}
-	}
 
-	// exec query
-	result := db.ExecQuery(ctx, test.Query)
-	assert.Empty(t, result.Errors, test.Description)
+		// exec query
+		result := db.ExecQuery(ctx, test.Query)
+		assert.Empty(t, result.Errors, test.Description)
 
-	resultantData := result.Data.([]map[string]interface{})
+		resultantData := result.Data.([]map[string]interface{})
 
-	fmt.Println(test.Description)
-	fmt.Println(result.Data)
-	fmt.Println("--------------")
-	fmt.Println("")
+		fmt.Println(test.Description)
+		fmt.Println(result.Data)
+		fmt.Println("--------------")
+		fmt.Println("")
 
-	// compare results
-	assert.Equal(t, len(test.Results), len(resultantData), test.Description)
-	for i, result := range resultantData {
-		assert.Equal(t, test.Results[i], result, test.Description)
+		// compare results
+		assert.Equal(t, len(test.Results), len(resultantData), test.Description)
+		for i, result := range resultantData {
+			assert.Equal(t, test.Results[i], result, test.Description)
+		}
 	}
 }
diff --git a/db/txn.go b/db/txn.go
@@ -70,22 +70,18 @@ func (db *DB) newTxn(ctx context.Context, readonly bool) (*Txn, error) {
 		txn.Txn = dstxn
 
 		// Note: db.rootstore now has type `ds.Batching`.
-	} else if batchStore := db.rootstore; ok { // we support Batching
-		batcher, err := batchStore.Batch(ctx)
+	} else {
+		batcher, err := db.rootstore.Batch(ctx)
 		if err != nil {
 			return nil, err
 		}
 
 		// hide a ds.Batching store as a ds.Txn
 		rb := shimBatcherTxn{
-			Read:  batchStore,
+			Read:  db.rootstore,
 			Batch: batcher,
 		}
 		txn.Txn = rb
-	} else {
-		// our datastore supports neither TxnDatastore or Batching
-		// for now return error
-		return nil, ErrNoTxnSupport
 	}
 
 	// add the wrapped datastores using the existing KeyTransform functions from the db