[Improvement] Use bulks in tests (#263)

Author: ajanikow

test/backup_test.go

@@ -869,18 +869,12 @@ func TestBackupRestoreWithViews(t *testing.T) {
 		go func(i int) {
 			defer wg.Done()
 
-			for j := 0; j < numDocs; j++ {
-
-				book := BookWithAuthor{
+			sendBulks(t, col, ctx, func(t *testing.T, j int) interface{} {
+				return BookWithAuthor{
 					Title:  fmt.Sprintf("Hello World - %d", j),
 					Author: fmt.Sprintf("Author - %d", i),
 				}
-
-				_, err := col.CreateDocument(ctx, book)
-				if err != nil {
-					t.Fatalf("Failed to create document %s", describe(err))
-				}
-			}
+			}, numDocs)
 		}(k)
 	}
 	wg.Wait()

test/clean.go

@@ -34,5 +34,9 @@ type remove interface {
 }
 
 func clean(t *testing.T, ctx context.Context, col remove) {
+	if col == nil {
+		return
+	}
+
 	require.NoError(t, col.Remove(ctx))
 }

test/cluster_test.go

@@ -264,10 +264,10 @@ func TestClusterMoveShard(t *testing.T) {
 		col, err := db.CreateCollection(ctx, "test_move_shard", &driver.CreateCollectionOptions{
 			NumberOfShards: 12,
 		})
-		defer clean(t, ctx, col)
 		if err != nil {
 			t.Fatalf("CreateCollection failed: %s", describe(err))
 		}
+		defer clean(t, ctx, col)
 		h, err := cl.Health(ctx)
 		if err != nil {
 			t.Fatalf("Health failed: %s", describe(err))
@@ -365,10 +365,10 @@ func TestClusterResignLeadership(t *testing.T) {
 			NumberOfShards:    12,
 			ReplicationFactor: 2,
 		})
-		defer clean(t, ctx, col)
 		if err != nil {
 			t.Fatalf("CreateCollection failed: %s", describe(err))
 		}
+		defer clean(t, ctx, col)
 		inv, err := cl.DatabaseInventory(ctx, db)
 		if err != nil {
 			t.Fatalf("DatabaseInventory failed: %s", describe(err))
@@ -453,21 +453,21 @@ func TestClusterMoveShardWithViews(t *testing.T) {
 		col, err := db.CreateCollection(ctx, "test_move_shard_with_view", &driver.CreateCollectionOptions{
 			NumberOfShards: 12,
 		})
-		clean(t, ctx, col)
 		if err != nil {
 			t.Fatalf("CreateCollection failed: %s", describe(err))
 		}
+		defer clean(t, ctx, col)
 		opts := &driver.ArangoSearchViewProperties{
 			Links: driver.ArangoSearchLinks{
 				"test_move_shard_with_view": driver.ArangoSearchElementProperties{},
 			},
 		}
 		viewName := "test_move_shard_view"
 		view, err := db.CreateArangoSearchView(ctx, viewName, opts)
-		clean(t, ctx, view)
 		if err != nil {
 			t.Fatalf("Failed to create view '%s': %s", viewName, describe(err))
 		}
+		defer clean(t, ctx, view)
 		h, err := cl.Health(ctx)
 		if err != nil {
 			t.Fatalf("Health failed: %s", describe(err))

test/cursor_test.go

@@ -348,12 +348,10 @@ func TestCreateStreamCursor(t *testing.T) {
 	}
 
 	// This might take a few seconds
-	for i := 0; i < 10000; i++ {
-		user := UserDoc{Name: "John", Age: i}
-		if _, err := col.CreateDocument(ctx, user); err != nil {
-			t.Fatalf("Expected success, got %s", describe(err))
-		}
-	}
+	docs := 10000
+	sendBulks(t, col, ctx, func(t *testing.T, i int) interface{} {
+		return UserDoc{Name: "John", Age: i}
+	}, docs)
 	t.Log("Completed inserting 10k docs")
 
 	const expectedResults int = 10 * 10000

test/server_mode_auth_test.go

@@ -62,7 +62,7 @@ func TestServerModeAndGrants(t *testing.T) {
 		db := ensureDatabase(ctx, c, "_system", nil, t)
 		colName := "server_mode_and_grants_test1"
 		col := ensureCollection(ctx, db, colName, nil, t)
-		clean(t, ctx, col)
+		defer clean(t, ctx, col)
 
 		// Get database & collection access
 		defaultDBAccess, err := u.GetDatabaseAccess(ctx, db)

test/util.go

@@ -23,6 +23,7 @@
 package test
 
 import (
+	"context"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
@@ -32,6 +33,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	driver "github.com/arangodb/go-driver"
 )
 
@@ -147,3 +150,41 @@ func retry(interval, timeout time.Duration, f func() error) error {
 		}
 	}
 }
+
+const bulkSize = 1000
+
+func sendBulks(t *testing.T, col driver.Collection, ctx context.Context, creator func(t *testing.T, i int) interface{}, size int) {
+	current := 0
+	t.Logf("Creating %d documents", size)
+
+	for {
+		t.Logf("Created %d/%d documents", current, size)
+		stepSize := min(bulkSize, size-current)
+		if stepSize == 0 {
+			return
+		}
+
+		objs := make([]interface{}, min(bulkSize, stepSize))
+		for i := 0; i < stepSize; i++ {
+			objs[i] = creator(t, current+i)
+		}
+
+		_, _, err := col.CreateDocuments(ctx, objs)
+		t.Logf("Creating %d documents", len(objs))
+		require.NoError(t, err)
+
+		current += stepSize
+	}
+}
+
+func min(max int, ints ...int) int {
+	z := max
+
+	for _, i := range ints {
+		if z > i {
+			z = i
+		}
+	}
+
+	return z
+}