Merge pull request #15 from idealo/unit_testing

feat: enhanced efficiency and predictability

.github/workflows/test.yml

@@ -0,0 +1,40 @@
+name: Test
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+    paths:
+      - '**.go'
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Setup Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: '>=1.22'
+          cache: true
+
+      - name: Cache Go modules
+        uses: actions/cache@v4
+        with:
+          path: ~/go/pkg/mod
+          key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
+          restore-keys: |
+            ${{ runner.os }}-go-
+
+      - name: Tests
+        run: |
+          go test -v ./...

.gitignore

@@ -1,2 +1,4 @@
 mongo-bench
 .idea
+*.csv
+*.bak

Makefile

@@ -1,13 +1,14 @@
 GOCMD=go
 GOBUILD=$(GOCMD) build
 GOCLEAN=$(GOCMD) clean
+GOTEST=$(GOCMD) test
 BINARY_NAME=mongo-bench
 MAIN_FILE=main.go
 
 .PHONY: all
 all: build
 
-build:
+build: test
 	$(GOBUILD) -o $(BINARY_NAME) $(MAIN_FILE)
 	@echo "Build complete: $(BINARY_NAME)"
 

go.mod

@@ -2,18 +2,25 @@ module benchmarking
 
 go 1.22
 
-require go.mongodb.org/mongo-driver v1.17.1
+require (
+	github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475
+	github.com/stretchr/testify v1.9.0
+	go.mongodb.org/mongo-driver v1.17.1
+)
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/klauspost/compress v1.13.6 // indirect
 	github.com/montanaflynn/stats v0.7.1 // indirect
-	github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/xdg-go/pbkdf2 v1.0.0 // indirect
 	github.com/xdg-go/scram v1.1.2 // indirect
 	github.com/xdg-go/stringprep v1.0.4 // indirect
 	github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect
 	golang.org/x/crypto v0.26.0 // indirect
 	golang.org/x/sync v0.8.0 // indirect
 	golang.org/x/text v0.17.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -8,8 +8,14 @@ github.com/klauspost/compress v1.13.6 h1:P76CopJELS0TiO2mebmnzgWaajssP/EszplttgQ
 github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
 github.com/montanaflynn/stats v0.7.1 h1:etflOAAHORrCC44V+aR6Ftzort912ZU+YLiSTuV8eaE=
 github.com/montanaflynn/stats v0.7.1/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475 h1:N/ElC8H3+5XpJzTSTfLsJV/mx9Q9g7kxmchpfZyxgzM=
 github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c=
 github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
 github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY=
@@ -50,3 +56,7 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -17,6 +17,73 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/options"
 )
 
+// CollectionAPI defines an interface for MongoDB operations, allowing for testing
+type CollectionAPI interface {
+	InsertOne(ctx context.Context, document interface{}) (*mongo.InsertOneResult, error)
+	UpdateOne(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error)
+	DeleteOne(ctx context.Context, filter interface{}) (*mongo.DeleteResult, error)
+	CountDocuments(ctx context.Context, filter interface{}) (int64, error)
+	Drop(ctx context.Context) error
+	Find(ctx context.Context, filter interface{}, opts ...*options.FindOptions) (*mongo.Cursor, error)
+}
+
+// MongoDBCollection is a wrapper around mongo.Collection to implement CollectionAPI
+type MongoDBCollection struct {
+	*mongo.Collection
+}
+
+func (c *MongoDBCollection) InsertOne(ctx context.Context, document interface{}) (*mongo.InsertOneResult, error) {
+	return c.Collection.InsertOne(ctx, document)
+}
+
+func (c *MongoDBCollection) UpdateOne(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
+	return c.Collection.UpdateOne(ctx, filter, update, opts...)
+}
+
+func (c *MongoDBCollection) DeleteOne(ctx context.Context, filter interface{}) (*mongo.DeleteResult, error) {
+	return c.Collection.DeleteOne(ctx, filter)
+}
+
+func (c *MongoDBCollection) CountDocuments(ctx context.Context, filter interface{}) (int64, error) {
+	return c.Collection.CountDocuments(ctx, filter)
+}
+
+func (c *MongoDBCollection) Drop(ctx context.Context) error {
+	return c.Collection.Drop(ctx)
+}
+
+func (c *MongoDBCollection) Find(ctx context.Context, filter interface{}, opts ...*options.FindOptions) (*mongo.Cursor, error) {
+	return c.Collection.Find(ctx, filter, opts...)
+}
+
+// fetchDocumentIDs fetches all document IDs from the collection for delete operations
+func fetchDocumentIDs(collection CollectionAPI) ([]int64, error) {
+	var docIDs []int64
+
+	cursor, err := collection.Find(context.Background(), bson.M{}, options.Find().SetProjection(bson.M{"_id": 1}))
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch document IDs: %v", err)
+	}
+	defer cursor.Close(context.Background())
+
+	for cursor.Next(context.Background()) {
+		var result bson.M
+		if err := cursor.Decode(&result); err != nil {
+			log.Printf("Failed to decode document: %v", err)
+			continue
+		}
+		if id, ok := result["_id"].(int64); ok {
+			docIDs = append(docIDs, id)
+		}
+	}
+
+	if err := cursor.Err(); err != nil {
+		return nil, fmt.Errorf("cursor error: %v", err)
+	}
+
+	return docIDs, nil
+}
+
 func main() {
 	var threads int
 	var docCount int
@@ -37,25 +104,24 @@ func main() {
 	}
 	defer client.Disconnect(context.Background())
 
-	// Run all tests in sequence if runAll flag is set
+	collection := client.Database("benchmarking").Collection("testdata")
+	mongoCollection := &MongoDBCollection{Collection: collection}
+
 	if runAll {
-		runTestSequence(client, threads, docCount)
+		runTestSequence(mongoCollection, threads, docCount)
 	} else {
-		// Run a single test based on testType
-		runTest(client, testType, threads, docCount)
+		runTest(mongoCollection, testType, threads, docCount, fetchDocumentIDs)
 	}
 }
 
-func runTestSequence(client *mongo.Client, threads, docCount int) {
+func runTestSequence(collection CollectionAPI, threads, docCount int) {
 	tests := []string{"insert", "update", "delete", "upsert"}
 	for _, test := range tests {
-		runTest(client, test, threads, docCount)
+		runTest(collection, test, threads, docCount, fetchDocumentIDs)
 	}
 }
 
-func runTest(client *mongo.Client, testType string, threads, docCount int) {
-	collection := client.Database("benchmarking").Collection("testdata")
-
+func runTest(collection CollectionAPI, testType string, threads, docCount int, fetchDocIDs func(CollectionAPI) ([]int64, error)) {
 	if testType == "insert" || testType == "upsert" {
 		if err := collection.Drop(context.Background()); err != nil {
 			log.Fatalf("Failed to drop collection: %v", err)
@@ -66,13 +132,39 @@ func runTest(client *mongo.Client, testType string, threads, docCount int) {
 	}
 
 	insertRate := metrics.NewMeter()
-
 	var records [][]string
 	records = append(records, []string{"t", "count", "mean", "m1_rate", "m5_rate", "m15_rate", "mean_rate"})
 
+	var partitions [][]int64 // To hold the document IDs or dummy IDs, partitioned for each thread
+
+	// Prepare partitions based on test type
+	switch testType {
+	case "delete":
+		// Fetch document IDs and partition them
+		docIDs, err := fetchDocIDs(collection)
+		if err != nil {
+			log.Fatalf("Failed to fetch document IDs: %v", err)
+		}
+		partitions = make([][]int64, threads)
+		for i, id := range docIDs {
+			partitions[i%threads] = append(partitions[i%threads], id)
+		}
+
+	case "insert", "update", "upsert":
+		// Generate unique or random IDs for insert/update/upsert
+		partitions = make([][]int64, threads)
+		for i := 0; i < docCount; i++ {
+			id := int64(i)
+			if testType == "update" || testType == "upsert" {
+				id = int64(rand.Intn(docCount)) // Random ID for update/upsert
+			}
+			partitions[i%threads] = append(partitions[i%threads], id)
+		}
+	}
+
+	// Start the ticker just before starting the main workload goroutines
 	secondTicker := time.NewTicker(1 * time.Second)
 	defer secondTicker.Stop()
-
 	go func() {
 		for range secondTicker.C {
 			timestamp := time.Now().Unix()
@@ -98,50 +190,17 @@ func runTest(client *mongo.Client, testType string, threads, docCount int) {
 		}
 	}()
 
-	var remainingDocIDs sync.Map
-
-	if testType == "delete" {
-		// Fetch all document IDs from the database to ensure they exist
-		cursor, err := collection.Find(context.Background(), bson.M{}, options.Find().SetProjection(bson.M{"_id": 1}))
-		if err != nil {
-			log.Fatalf("Failed to fetch document IDs: %v", err)
-		}
-		defer cursor.Close(context.Background())
-
-		for cursor.Next(context.Background()) {
-			var result bson.M
-			if err := cursor.Decode(&result); err != nil {
-				log.Printf("Failed to decode document: %v", err)
-				continue
-			}
-			if id, ok := result["_id"].(int64); ok {
-				remainingDocIDs.Store(id, true)
-			}
-		}
-
-		if err := cursor.Err(); err != nil {
-			log.Fatalf("Cursor error: %v", err)
-		}
-	}
-
+	// Launch threads based on the specific workload type
 	var wg sync.WaitGroup
 	wg.Add(threads)
 
 	for i := 0; i < threads; i++ {
-		go func(threadID int) {
+		go func(partition []int64) {
 			defer wg.Done()
-			threadDocCount := docCount / threads
-			for j := 0; j < threadDocCount; j++ {
+			for _, docID := range partition {
 				switch testType {
 				case "insert":
-					docID := int64(threadID*threadDocCount + j)
-					doc := bson.M{
-						"_id":            docID,
-						"threadId":       threadID,
-						"threadRunCount": 1,
-						"rnd":            rand.Int63(),
-						"v":              1,
-					}
+					doc := bson.M{"_id": docID, "threadRunCount": 1, "rnd": rand.Int63(), "v": 1}
 					_, err := collection.InsertOne(context.Background(), doc)
 					if err == nil {
 						insertRate.Mark(1)
@@ -150,7 +209,6 @@ func runTest(client *mongo.Client, testType string, threads, docCount int) {
 					}
 
 				case "update":
-					docID := int64(rand.Intn(docCount))
 					filter := bson.M{"_id": docID}
 					update := bson.M{"$set": bson.M{"updatedAt": time.Now().Unix(), "rnd": rand.Int63()}}
 					_, err := collection.UpdateOne(context.Background(), filter, update)
@@ -161,7 +219,6 @@ func runTest(client *mongo.Client, testType string, threads, docCount int) {
 					}
 
 				case "upsert":
-					docID := int64(rand.Intn(docCount / 2))
 					filter := bson.M{"_id": docID}
 					update := bson.M{"$set": bson.M{"updatedAt": time.Now().Unix(), "rnd": rand.Int63()}}
 					opts := options.Update().SetUpsert(true)
@@ -173,38 +230,23 @@ func runTest(client *mongo.Client, testType string, threads, docCount int) {
 					}
 
 				case "delete":
-					for {
-						var docID int64
-						found := false
-
-						remainingDocIDs.Range(func(key, value interface{}) bool {
-							docID = key.(int64)
-							found = true
-							return false
-						})
-
-						if !found {
-							log.Println("No documents left to delete.")
-							return
-						}
-
-						filter := bson.M{"_id": docID}
-						result, err := collection.DeleteOne(context.Background(), filter)
-						if err != nil {
-							log.Printf("Delete failed: %v", err)
-							break
-						} else if result.DeletedCount > 0 {
-							insertRate.Mark(1)
-							remainingDocIDs.Delete(docID)
-						}
+					filter := bson.M{"_id": docID}
+					result, err := collection.DeleteOne(context.Background(), filter)
+					if err != nil {
+						log.Printf("Delete failed for _id %d: %v", docID, err)
+						continue // Move to next document without retrying
+					}
+					if result.DeletedCount > 0 {
+						insertRate.Mark(1)
 					}
 				}
 			}
-		}(i)
+		}(partitions[i])
 	}
 
 	wg.Wait()
 
+	// Final metrics recording
 	timestamp := time.Now().Unix()
 	count := insertRate.Count()
 	mean := insertRate.RateMean()