90 lines
2.3 KiB
Go
90 lines
2.3 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)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// ExecutePipeline 执行管道（用于测试）
|
|
func (e *AggregationEngine) ExecutePipeline(docs []types.Document, pipeline []types.AggregateStage) ([]types.Document, error) {
|
|
result := docs
|
|
for _, stage := range pipeline {
|
|
var err error
|
|
result, err = e.executeStage(stage, result)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|