github.com/crowdsecurity/crowdsec@v1.6.1/pkg/csplugin/watcher_test.go (about)

     1  package csplugin
     2  
     3  import (
     4  	"context"
     5  	"log"
     6  	"runtime"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/require"
    11  	"gopkg.in/tomb.v2"
    12  
    13  	"github.com/crowdsecurity/go-cs-lib/cstest"
    14  
    15  	"github.com/crowdsecurity/crowdsec/pkg/models"
    16  )
    17  
    18  var ctx = context.Background()
    19  
    20  func resetTestTomb(testTomb *tomb.Tomb, pw *PluginWatcher) {
    21  	testTomb.Kill(nil)
    22  	<-pw.PluginEvents
    23  	if err := testTomb.Wait(); err != nil {
    24  		log.Fatal(err)
    25  	}
    26  }
    27  
    28  func resetWatcherAlertCounter(pw *PluginWatcher) {
    29  	pw.AlertCountByPluginName.Lock()
    30  	for k := range pw.AlertCountByPluginName.data {
    31  		pw.AlertCountByPluginName.data[k] = 0
    32  	}
    33  	pw.AlertCountByPluginName.Unlock()
    34  }
    35  
    36  func insertNAlertsToPlugin(pw *PluginWatcher, n int, pluginName string) {
    37  	for i := 0; i < n; i++ {
    38  		pw.Inserts <- pluginName
    39  	}
    40  }
    41  
    42  func listenChannelWithTimeout(ctx context.Context, channel chan string) error {
    43  	select {
    44  	case x := <-channel:
    45  		log.Printf("received -> %v", x)
    46  	case <-ctx.Done():
    47  		return ctx.Err()
    48  	}
    49  	return nil
    50  }
    51  
    52  func TestPluginWatcherInterval(t *testing.T) {
    53  	if runtime.GOOS == "windows" {
    54  		t.Skip("Skipping test on windows because timing is not reliable")
    55  	}
    56  	pw := PluginWatcher{}
    57  	alertsByPluginName := make(map[string][]*models.Alert)
    58  	testTomb := tomb.Tomb{}
    59  	configs := map[string]PluginConfig{
    60  		"testPlugin": {
    61  			GroupWait: time.Millisecond,
    62  		},
    63  	}
    64  	pw.Init(configs, alertsByPluginName)
    65  	pw.Start(&testTomb)
    66  
    67  	ct, cancel := context.WithTimeout(ctx, time.Microsecond)
    68  	defer cancel()
    69  	err := listenChannelWithTimeout(ct, pw.PluginEvents)
    70  	cstest.RequireErrorContains(t, err, "context deadline exceeded")
    71  	resetTestTomb(&testTomb, &pw)
    72  	testTomb = tomb.Tomb{}
    73  	pw.Start(&testTomb)
    74  
    75  	ct, cancel = context.WithTimeout(ctx, time.Millisecond*5)
    76  	defer cancel()
    77  	err = listenChannelWithTimeout(ct, pw.PluginEvents)
    78  	require.NoError(t, err)
    79  	resetTestTomb(&testTomb, &pw)
    80  	// This is to avoid the int complaining
    81  }
    82  
    83  func TestPluginAlertCountWatcher(t *testing.T) {
    84  	if runtime.GOOS == "windows" {
    85  		t.Skip("Skipping test on windows because timing is not reliable")
    86  	}
    87  	pw := PluginWatcher{}
    88  	alertsByPluginName := make(map[string][]*models.Alert)
    89  	configs := map[string]PluginConfig{
    90  		"testPlugin": {
    91  			GroupThreshold: 5,
    92  		},
    93  	}
    94  	testTomb := tomb.Tomb{}
    95  	pw.Init(configs, alertsByPluginName)
    96  	pw.Start(&testTomb)
    97  
    98  	// Channel won't contain any events since threshold is not crossed.
    99  	ct, cancel := context.WithTimeout(ctx, time.Second)
   100  	defer cancel()
   101  	err := listenChannelWithTimeout(ct, pw.PluginEvents)
   102  	cstest.RequireErrorContains(t, err, "context deadline exceeded")
   103  
   104  	// Channel won't contain any events since threshold is not crossed.
   105  	resetWatcherAlertCounter(&pw)
   106  	insertNAlertsToPlugin(&pw, 4, "testPlugin")
   107  	ct, cancel = context.WithTimeout(ctx, time.Second)
   108  	defer cancel()
   109  	err = listenChannelWithTimeout(ct, pw.PluginEvents)
   110  	cstest.RequireErrorContains(t, err, "context deadline exceeded")
   111  
   112  	// Channel will contain an event since threshold is crossed.
   113  	resetWatcherAlertCounter(&pw)
   114  	insertNAlertsToPlugin(&pw, 5, "testPlugin")
   115  	ct, cancel = context.WithTimeout(ctx, time.Second)
   116  	defer cancel()
   117  	err = listenChannelWithTimeout(ct, pw.PluginEvents)
   118  	require.NoError(t, err)
   119  	resetTestTomb(&testTomb, &pw)
   120  }