github.com/crowdsecurity/crowdsec@v1.6.1/pkg/leakybucket/manager_run_test.go (about)

     1  package leakybucket
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/crowdsecurity/crowdsec/pkg/types"
     9  	log "github.com/sirupsen/logrus"
    10  	"gopkg.in/tomb.v2"
    11  )
    12  
    13  func expectBucketCount(buckets *Buckets, expected int) error {
    14  	count := 0
    15  	buckets.Bucket_map.Range(func(rkey, rvalue interface{}) bool {
    16  		count++
    17  		return true
    18  	})
    19  	if count != expected {
    20  		return fmt.Errorf("expected %d live buckets, got %d", expected, count)
    21  	}
    22  	return nil
    23  
    24  }
    25  
    26  func TestGCandDump(t *testing.T) {
    27  	var (
    28  		buckets = NewBuckets()
    29  		tomb    = &tomb.Tomb{}
    30  	)
    31  
    32  	var Holders = []BucketFactory{
    33  		//one overflowing soon + bh
    34  		{
    35  			Name:        "test_counter_fast",
    36  			Description: "test_counter_fast",
    37  			Debug:       true,
    38  			Type:        "counter",
    39  			Capacity:    -1,
    40  			Duration:    "0.5s",
    41  			Blackhole:   "1m",
    42  			Filter:      "true",
    43  			wgDumpState: buckets.wgDumpState,
    44  			wgPour:      buckets.wgPour,
    45  		},
    46  		//one long counter
    47  		{
    48  			Name:        "test_counter_slow",
    49  			Description: "test_counter_slow",
    50  			Debug:       true,
    51  			Type:        "counter",
    52  			Capacity:    -1,
    53  			Duration:    "10m",
    54  			Filter:      "true",
    55  			wgDumpState: buckets.wgDumpState,
    56  			wgPour:      buckets.wgPour,
    57  		},
    58  		//slow leaky
    59  		{
    60  			Name:        "test_leaky_slow",
    61  			Description: "test_leaky_slow",
    62  			Debug:       true,
    63  			Type:        "leaky",
    64  			Capacity:    5,
    65  			LeakSpeed:   "10m",
    66  			Filter:      "true",
    67  			wgDumpState: buckets.wgDumpState,
    68  			wgPour:      buckets.wgPour,
    69  		},
    70  	}
    71  
    72  	for idx := range Holders {
    73  		if err := LoadBucket(&Holders[idx], tomb); err != nil {
    74  			t.Fatalf("while loading (%d/%d): %s", idx, len(Holders), err)
    75  		}
    76  		if err := ValidateFactory(&Holders[idx]); err != nil {
    77  			t.Fatalf("while validating (%d/%d): %s", idx, len(Holders), err)
    78  		}
    79  	}
    80  
    81  	log.Printf("Pouring to bucket")
    82  
    83  	var in = types.Event{Parsed: map[string]string{"something": "something"}}
    84  	//pour an item that will go to leaky + counter
    85  	ok, err := PourItemToHolders(in, Holders, buckets)
    86  	if err != nil {
    87  		t.Fatalf("while pouring item : %s", err)
    88  	}
    89  	if !ok {
    90  		t.Fatalf("didn't pour item")
    91  	}
    92  
    93  	time.Sleep(2 * time.Second)
    94  
    95  	if err := expectBucketCount(buckets, 3); err != nil {
    96  		t.Fatal(err)
    97  	}
    98  	log.Printf("Bucket GC")
    99  
   100  	//call garbage collector
   101  	if err := GarbageCollectBuckets(time.Now().UTC(), buckets); err != nil {
   102  		t.Fatalf("failed to garbage collect buckets : %s", err)
   103  	}
   104  
   105  	if err := expectBucketCount(buckets, 1); err != nil {
   106  		t.Fatal(err)
   107  	}
   108  
   109  	log.Printf("Dumping buckets state")
   110  	//dump remaining buckets
   111  	if _, err := DumpBucketsStateAt(time.Now().UTC(), ".", buckets); err != nil {
   112  		t.Fatalf("failed to dump buckets : %s", err)
   113  	}
   114  }
   115  
   116  func TestShutdownBuckets(t *testing.T) {
   117  	var (
   118  		buckets = NewBuckets()
   119  		Holders = []BucketFactory{
   120  			//one long counter
   121  			{
   122  				Name:        "test_counter_slow",
   123  				Description: "test_counter_slow",
   124  				Debug:       true,
   125  				Type:        "counter",
   126  				Capacity:    -1,
   127  				Duration:    "10m",
   128  				Filter:      "true",
   129  				wgDumpState: buckets.wgDumpState,
   130  				wgPour:      buckets.wgPour,
   131  			},
   132  			//slow leaky
   133  			{
   134  				Name:        "test_leaky_slow",
   135  				Description: "test_leaky_slow",
   136  				Debug:       true,
   137  				Type:        "leaky",
   138  				Capacity:    5,
   139  				LeakSpeed:   "10m",
   140  				Filter:      "true",
   141  				wgDumpState: buckets.wgDumpState,
   142  				wgPour:      buckets.wgPour,
   143  			},
   144  		}
   145  		tomb = &tomb.Tomb{}
   146  	)
   147  
   148  	for idx := range Holders {
   149  		if err := LoadBucket(&Holders[idx], tomb); err != nil {
   150  			t.Fatalf("while loading (%d/%d): %s", idx, len(Holders), err)
   151  		}
   152  		if err := ValidateFactory(&Holders[idx]); err != nil {
   153  			t.Fatalf("while validating (%d/%d): %s", idx, len(Holders), err)
   154  		}
   155  	}
   156  
   157  	log.Printf("Pouring to bucket")
   158  
   159  	var in = types.Event{Parsed: map[string]string{"something": "something"}}
   160  	//pour an item that will go to leaky + counter
   161  	ok, err := PourItemToHolders(in, Holders, buckets)
   162  	if err != nil {
   163  		t.Fatalf("while pouring item : %s", err)
   164  	}
   165  	if !ok {
   166  		t.Fatalf("didn't pour item")
   167  	}
   168  
   169  	time.Sleep(1 * time.Second)
   170  
   171  	if err := expectBucketCount(buckets, 2); err != nil {
   172  		t.Fatal(err)
   173  	}
   174  	if err := ShutdownAllBuckets(buckets); err != nil {
   175  		t.Fatalf("while shutting down buckets : %s", err)
   176  	}
   177  	time.Sleep(2 * time.Second)
   178  	if err := expectBucketCount(buckets, 2); err != nil {
   179  		t.Fatal(err)
   180  	}
   181  
   182  }