generate data

FerretDB · AlekSi · Apr 18, 2023 · Apr 5, 2023 · Apr 5, 2023 · Apr 5, 2023
commit 8955cbc342fd067b6a774fa757036c8cc4737123
@@ -29,7 +29,7 @@ type QueryBenchmarkCase struct {
 }
 
 func BenchmarkFoo(b *testing.B) {
-	ctx, coll, collNoPushdown, _ := setup.SetupBenchmark(b)
+	ctx, coll, collNoPushdown, compatColl := setup.SetupBenchmark(b, setup.SimpleData)
 
 	for name, bm := range map[string]QueryBenchmarkCase{
 		"String": {
@@ -55,7 +55,15 @@ func BenchmarkFoo(b *testing.B) {
 					require.NoError(b, cur.All(ctx, &res))
 				}
 			})
+			b.Run("Compat", func(b *testing.B) {
+				for i := 0; i < b.N; i++ {
+					cur, err := compatColl.Find(ctx, bm.filter)
+					require.NoError(b, err)
 
+					var res []bson.D
+					require.NoError(b, cur.All(ctx, &res))
+				}
+			})
 		})
 	}
 }

@@ -155,15 +155,27 @@ func SetupWithOpts(tb testing.TB, opts *SetupOpts) *SetupResult {
 	}
 }
 
-func SetupBenchmark(tb testing.TB) (context.Context, *mongo.Collection, *mongo.Collection, *mongo.Collection) {
+type BenchmarkData func(context.Context, *mongo.Collection) error
+
+var SimpleData BenchmarkData = func(ctx context.Context, coll *mongo.Collection) error {
+	var docs []any
+	for i, doc := range []any{
+		"foo", 42, "42", bson.D{{"42", "hello"}},
+	} {
+		docs = append(docs, bson.D{{"_id", i}, {"v", doc}})
+	}
+
+	_, err := coll.InsertMany(ctx, docs)
+	return err
+}
+
+func SetupBenchmark(tb testing.TB, insert BenchmarkData) (context.Context, *mongo.Collection, *mongo.Collection, *mongo.Collection) {
 	tb.Helper()
 
 	if *compatURLF == "" {
 		tb.Skip("-compat-url is empty, skipping benchmark")
 	}
 
-	// TODO insert data
-
 	ctx, cancel := context.WithCancel(testutil.Ctx(tb))
 
 	level := zap.NewAtomicLevelAt(zap.ErrorLevel)
@@ -177,17 +189,22 @@ func SetupBenchmark(tb testing.TB) (context.Context, *mongo.Collection, *mongo.C
 
 	compatClient := setupClient(tb, ctx, *compatURLF)
 
-	// TODO "compat" collection
-
 	tb.Cleanup(cancel)
 
-	dbName := tb.Name()
-	collName := "Benchmark"
+	//dbName := tb.Name()
+	//collName := "Benchmark"
+
+	var collections []*mongo.Collection
+
+	for _, client := range []*mongo.Client{
+		targetClient, targetNoPushdownClient, compatClient,
+	} {
+		coll := setupCollection(tb, ctx, client, &SetupOpts{})
+		require.NoError(tb, insert(ctx, coll))
+		collections = append(collections, coll)
+	}
 
-	return ctx,
-		targetClient.Database(dbName).Collection(collName),
-		targetNoPushdownClient.Database(dbName).Collection(collName),
-		compatClient.Database(dbName).Collection(collName)
+	return ctx, collections[0], collections[1], collections[2]
 }
 
 // Setup setups a single collection for all compatible providers, if the are present.