gomog/internal/engine/aggregate_test.go

77 lines
1.9 KiB
Go

package engine
import (
"testing"
"git.kingecg.top/kingecg/gomog/pkg/types"
)
func TestAggregationEngine(t *testing.T) {
store := NewMemoryStore(nil)
engine := NewAggregationEngine(store)
// 创建测试数据
docs := []types.Document{
{ID: "1", Data: map[string]interface{}{"category": "A", "amount": 100, "status": "completed"}},
{ID: "2", Data: map[string]interface{}{"category": "A", "amount": 150, "status": "completed"}},
{ID: "3", Data: map[string]interface{}{"category": "B", "amount": 200, "status": "pending"}},
{ID: "4", Data: map[string]interface{}{"category": "B", "amount": 50, "status": "completed"}},
}
tests := []struct {
name string
pipeline []types.AggregateStage
wantLen int
}{
{
name: "match stage",
pipeline: []types.AggregateStage{
{Stage: "$match", Spec: map[string]interface{}{"status": "completed"}},
},
wantLen: 3,
},
{
name: "group stage",
pipeline: []types.AggregateStage{
{Stage: "$group", Spec: map[string]interface{}{
"_id": "$category",
"total": map[string]interface{}{"$sum": "$amount"},
}},
},
wantLen: 2,
},
{
name: "match and group",
pipeline: []types.AggregateStage{
{Stage: "$match", Spec: map[string]interface{}{"status": "completed"}},
{Stage: "$group", Spec: map[string]interface{}{
"_id": "$category",
"total": map[string]interface{}{"$sum": "$amount"},
}},
},
wantLen: 2,
},
{
name: "sort and limit",
pipeline: []types.AggregateStage{
{Stage: "$sort", Spec: map[string]interface{}{"amount": -1}},
{Stage: "$limit", Spec: 2},
},
wantLen: 2,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result, err := engine.ExecutePipeline(docs, tt.pipeline)
if err != nil {
t.Errorf("ExecutePipeline() error = %v", err)
return
}
if len(result) != tt.wantLen {
t.Errorf("ExecutePipeline() result length = %d, want %d", len(result), tt.wantLen)
}
})
}
}