367 lines
9.3 KiB
Go
367 lines
9.3 KiB
Go
package engine
|
||
|
||
import (
|
||
"fmt"
|
||
"testing"
|
||
|
||
"git.kingecg.top/kingecg/gomog/pkg/types"
|
||
)
|
||
|
||
// ========== 辅助函数：生成测试数据 ==========
|
||
|
||
func generateDocuments(count int) map[string]types.Document {
|
||
docs := make(map[string]types.Document)
|
||
for i := 0; i < count; i++ {
|
||
docs[fmt.Sprintf("doc%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("doc%d", i),
|
||
Data: map[string]interface{}{
|
||
"name": fmt.Sprintf("Item%d", i),
|
||
"value": float64(i),
|
||
"category": fmt.Sprintf("cat%d", i%10),
|
||
"status": map[string]interface{}{"active": true, "priority": float64(i % 5)},
|
||
},
|
||
}
|
||
}
|
||
return docs
|
||
}
|
||
|
||
// ========== 聚合管道基准测试 ==========
|
||
|
||
// BenchmarkAggregationPipeline_Simple 简单聚合管道性能测试
|
||
func BenchmarkAggregationPipeline_Simple(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
// 准备 100 个文档
|
||
CreateTestCollectionForTesting(store, "benchmark_simple", generateDocuments(100))
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$match", Spec: map[string]interface{}{"status.active": true}},
|
||
{Stage: "$limit", Spec: float64(10)},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("benchmark_simple", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkAggregationPipeline_Group 分组聚合性能测试
|
||
func BenchmarkAggregationPipeline_Group(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
// 生成 1000 个文档
|
||
docs := make(map[string]types.Document)
|
||
for i := 0; i < 1000; i++ {
|
||
docs[fmt.Sprintf("doc%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("doc%d", i),
|
||
Data: map[string]interface{}{
|
||
"category": fmt.Sprintf("cat%d", i%10), // 10 个类别
|
||
"value": float64(i),
|
||
},
|
||
}
|
||
}
|
||
CreateTestCollectionForTesting(store, "benchmark_group", docs)
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$group", Spec: map[string]interface{}{
|
||
"_id": "$category",
|
||
"total": map[string]interface{}{"$sum": "$value"},
|
||
"count": map[string]interface{}{"$sum": float64(1)},
|
||
}},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("benchmark_group", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkAggregationPipeline_Complex 复杂聚合管道性能测试
|
||
func BenchmarkAggregationPipeline_Complex(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
// 主集合：500 个订单
|
||
mainDocs := make(map[string]types.Document)
|
||
for i := 0; i < 500; i++ {
|
||
mainDocs[fmt.Sprintf("main%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("main%d", i),
|
||
Data: map[string]interface{}{
|
||
"user_id": float64(i % 100),
|
||
"amount": float64(i * 10),
|
||
"status": "completed",
|
||
},
|
||
}
|
||
}
|
||
CreateTestCollectionForTesting(store, "orders", mainDocs)
|
||
|
||
// 关联集合：100 个用户
|
||
userDocs := make(map[string]types.Document)
|
||
for i := 0; i < 100; i++ {
|
||
userDocs[fmt.Sprintf("user%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("user%d", i),
|
||
Data: map[string]interface{}{
|
||
"_id": float64(i),
|
||
"name": fmt.Sprintf("User%d", i),
|
||
"department": fmt.Sprintf("Dept%d", i%5),
|
||
},
|
||
}
|
||
}
|
||
CreateTestCollectionForTesting(store, "users", userDocs)
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$match", Spec: map[string]interface{}{"status": "completed"}},
|
||
{Stage: "$lookup", Spec: map[string]interface{}{
|
||
"from": "users",
|
||
"localField": "user_id",
|
||
"foreignField": "_id",
|
||
"as": "user_info",
|
||
}},
|
||
{Stage: "$unwind", Spec: "$user_info"},
|
||
{Stage: "$group", Spec: map[string]interface{}{
|
||
"_id": "$user_info.department",
|
||
"total_sales": map[string]interface{}{"$sum": "$amount"},
|
||
}},
|
||
{Stage: "$sort", Spec: map[string]interface{}{"total_sales": -1}},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("orders", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ========== 查询操作符基准测试 ==========
|
||
|
||
// BenchmarkQuery_Expr 表达式查询性能测试
|
||
func BenchmarkQuery_Expr(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
docs := make(map[string]types.Document)
|
||
for i := 0; i < 1000; i++ {
|
||
docs[fmt.Sprintf("doc%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("doc%d", i),
|
||
Data: map[string]interface{}{
|
||
"score": float64(i),
|
||
"name": fmt.Sprintf("item%d", i),
|
||
},
|
||
}
|
||
}
|
||
CreateTestCollectionForTesting(store, "benchmark_expr", docs)
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$match", Spec: map[string]interface{}{
|
||
"$expr": map[string]interface{}{
|
||
"$gt": []interface{}{"$score", float64(500)},
|
||
},
|
||
}},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("benchmark_expr", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkQuery_JsonSchema JSON Schema 验证性能测试
|
||
func BenchmarkQuery_JsonSchema(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
docs := make(map[string]types.Document)
|
||
for i := 0; i < 500; i++ {
|
||
docs[fmt.Sprintf("doc%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("doc%d", i),
|
||
Data: map[string]interface{}{
|
||
"name": fmt.Sprintf("item%d", i),
|
||
"price": float64(i * 10),
|
||
"stock": float64(i),
|
||
},
|
||
}
|
||
}
|
||
CreateTestCollectionForTesting(store, "benchmark_schema", docs)
|
||
|
||
schema := map[string]interface{}{
|
||
"properties": map[string]interface{}{
|
||
"price": map[string]interface{}{
|
||
"bsonType": "number",
|
||
"minimum": float64(100),
|
||
},
|
||
},
|
||
}
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$match", Spec: map[string]interface{}{
|
||
"$jsonSchema": schema,
|
||
}},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("benchmark_schema", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ========== 类型转换基准测试 ==========
|
||
|
||
// BenchmarkTypeConversion_ToString 字符串转换性能测试
|
||
func BenchmarkTypeConversion_ToString(b *testing.B) {
|
||
engine := &AggregationEngine{}
|
||
data := map[string]interface{}{"value": float64(12345)}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_ = engine.toString("$value", data)
|
||
}
|
||
}
|
||
|
||
// BenchmarkTypeConversion_Bitwise 位运算性能测试
|
||
func BenchmarkTypeConversion_Bitwise(b *testing.B) {
|
||
engine := &AggregationEngine{}
|
||
operand := []interface{}{float64(12345), float64(67890)}
|
||
data := map[string]interface{}{}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_ = engine.bitAnd(operand, data)
|
||
}
|
||
}
|
||
|
||
// ========== 投影基准测试 ==========
|
||
|
||
// BenchmarkProjection_ElemMatch 数组元素匹配性能测试
|
||
func BenchmarkProjection_ElemMatch(b *testing.B) {
|
||
data := map[string]interface{}{
|
||
"scores": []interface{}{
|
||
map[string]interface{}{"subject": "math", "score": float64(85)},
|
||
map[string]interface{}{"subject": "english", "score": float64(92)},
|
||
map[string]interface{}{"subject": "science", "score": float64(78)},
|
||
},
|
||
}
|
||
spec := map[string]interface{}{
|
||
"$elemMatch": map[string]interface{}{
|
||
"score": map[string]interface{}{"$gte": float64(90)},
|
||
},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_ = projectElemMatch(data, "scores", spec)
|
||
}
|
||
}
|
||
|
||
// BenchmarkProjection_Slice 数组切片性能测试
|
||
func BenchmarkProjection_Slice(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
docs := make(map[string]types.Document)
|
||
for i := 0; i < 100; i++ {
|
||
docs[fmt.Sprintf("doc%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("doc%d", i),
|
||
Data: map[string]interface{}{
|
||
"tags": []interface{}{"tag1", "tag2", "tag3", "tag4", "tag5"},
|
||
},
|
||
}
|
||
}
|
||
CreateTestCollectionForTesting(store, "slice_bench", docs)
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$project", Spec: map[string]interface{}{
|
||
"tags": map[string]interface{}{"$slice": float64(3)},
|
||
}},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("slice_bench", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ========== UnionWith 基准测试 ==========
|
||
|
||
// BenchmarkUnionWith_Simple 集合并集性能测试（无 pipeline）
|
||
func BenchmarkUnionWith_Simple(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
CreateTestCollectionForTesting(store, "union_main", generateDocuments(100))
|
||
CreateTestCollectionForTesting(store, "union_other", generateDocuments(100))
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$unionWith", Spec: "union_other"},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("union_main", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ========== Redact 基准测试 ==========
|
||
|
||
// BenchmarkRedact_LevelBased 基于层级的文档红黑性能测试
|
||
func BenchmarkRedact_LevelBased(b *testing.B) {
|
||
store := NewMemoryStore(nil)
|
||
engine := NewAggregationEngine(store)
|
||
|
||
docs := make(map[string]types.Document)
|
||
for i := 0; i < 200; i++ {
|
||
docs[fmt.Sprintf("doc%d", i)] = types.Document{
|
||
ID: fmt.Sprintf("doc%d", i),
|
||
Data: map[string]interface{}{
|
||
"level": float64(i % 10),
|
||
"secret": "classified",
|
||
"public": "visible",
|
||
},
|
||
}
|
||
}
|
||
CreateTestCollectionForTesting(store, "redact_bench", docs)
|
||
|
||
spec := map[string]interface{}{
|
||
"$cond": map[string]interface{}{
|
||
"if": map[string]interface{}{
|
||
"$gte": []interface{}{"$level", float64(5)},
|
||
},
|
||
"then": "$$KEEP",
|
||
"else": "$$PRUNE",
|
||
},
|
||
}
|
||
|
||
pipeline := []types.AggregateStage{
|
||
{Stage: "$redact", Spec: spec},
|
||
}
|
||
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
_, err := engine.Execute("redact_bench", pipeline)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|