125 lines
3.5 KiB
Go
125 lines
3.5 KiB
Go
package leakybucket
|
|
|
|
import (
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/crowdsecurity/crowdsec/pkg/types"
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
func expectBucketCount(buckets *Buckets, expected int) error {
|
|
count := 0
|
|
buckets.Bucket_map.Range(func(rkey, rvalue interface{}) bool {
|
|
count++
|
|
return true
|
|
})
|
|
if count != expected {
|
|
return fmt.Errorf("expected %d live buckets, got %d", expected, count)
|
|
}
|
|
return nil
|
|
|
|
}
|
|
|
|
func TestGCandDump(t *testing.T) {
|
|
var buckets *Buckets = NewBuckets()
|
|
|
|
var Holders = []BucketFactory{
|
|
//one overflowing soon + bh
|
|
BucketFactory{Name: "test_counter_fast", Description: "test_counter_fast", Debug: true, Type: "counter", Capacity: -1, Duration: "0.5s", Blackhole: "1m", Filter: "true"},
|
|
//one long counter
|
|
BucketFactory{Name: "test_counter_slow", Description: "test_counter_slow", Debug: true, Type: "counter", Capacity: -1, Duration: "10m", Filter: "true"},
|
|
//slow leaky
|
|
BucketFactory{Name: "test_leaky_slow", Description: "test_leaky_slow", Debug: true, Type: "leaky", Capacity: 5, LeakSpeed: "10m", Filter: "true"},
|
|
}
|
|
|
|
for idx := range Holders {
|
|
if err := LoadBucket(&Holders[idx]); err != nil {
|
|
t.Fatalf("while loading (%d/%d): %s", idx, len(Holders), err)
|
|
}
|
|
if err := ValidateFactory(&Holders[idx]); err != nil {
|
|
t.Fatalf("while validating (%d/%d): %s", idx, len(Holders), err)
|
|
}
|
|
}
|
|
|
|
log.Printf("Pouring to bucket")
|
|
|
|
var in = types.Event{Parsed: map[string]string{"something": "something"}}
|
|
//pour an item that will go to leaky + counter
|
|
ok, err := PourItemToHolders(in, Holders, buckets)
|
|
if err != nil {
|
|
t.Fatalf("while pouring item : %s", err)
|
|
}
|
|
if !ok {
|
|
t.Fatalf("didn't pour item")
|
|
}
|
|
|
|
time.Sleep(2 * time.Second)
|
|
|
|
if err := expectBucketCount(buckets, 3); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
log.Printf("Bucket GC")
|
|
|
|
//call garbage collector
|
|
if err := GarbageCollectBuckets(time.Now(), buckets); err != nil {
|
|
t.Fatalf("failed to garbage collect buckets : %s", err)
|
|
}
|
|
|
|
if err := expectBucketCount(buckets, 1); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
log.Printf("Dumping buckets state")
|
|
//dump remaining buckets
|
|
if _, err := DumpBucketsStateAt(time.Now(), ".", buckets); err != nil {
|
|
t.Fatalf("failed to dump buckets : %s", err)
|
|
}
|
|
}
|
|
|
|
func TestShutdownBuckets(t *testing.T) {
|
|
var buckets *Buckets = NewBuckets()
|
|
var Holders = []BucketFactory{
|
|
//one long counter
|
|
BucketFactory{Name: "test_counter_slow", Description: "test_counter_slow", Debug: true, Type: "counter", Capacity: -1, Duration: "10m", Filter: "true"},
|
|
//slow leaky
|
|
BucketFactory{Name: "test_leaky_slow", Description: "test_leaky_slow", Debug: true, Type: "leaky", Capacity: 5, LeakSpeed: "10m", Filter: "true"},
|
|
}
|
|
|
|
for idx := range Holders {
|
|
if err := LoadBucket(&Holders[idx]); err != nil {
|
|
t.Fatalf("while loading (%d/%d): %s", idx, len(Holders), err)
|
|
}
|
|
if err := ValidateFactory(&Holders[idx]); err != nil {
|
|
t.Fatalf("while validating (%d/%d): %s", idx, len(Holders), err)
|
|
}
|
|
}
|
|
|
|
log.Printf("Pouring to bucket")
|
|
|
|
var in = types.Event{Parsed: map[string]string{"something": "something"}}
|
|
//pour an item that will go to leaky + counter
|
|
ok, err := PourItemToHolders(in, Holders, buckets)
|
|
if err != nil {
|
|
t.Fatalf("while pouring item : %s", err)
|
|
}
|
|
if !ok {
|
|
t.Fatalf("didn't pour item")
|
|
}
|
|
|
|
time.Sleep(1 * time.Second)
|
|
|
|
if err := expectBucketCount(buckets, 2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := ShutdownAllBuckets(buckets); err != nil {
|
|
t.Fatalf("while shuting down buckets : %s", err)
|
|
}
|
|
time.Sleep(2 * time.Second)
|
|
if err := expectBucketCount(buckets, 2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
}
|