github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/agent/daemon/enrichment/service_test.go (about)

     1  package enrichment
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	castpb "github.com/castai/kvisor/api/v1/runtime"
     9  	"github.com/castai/kvisor/pkg/ebpftracer/types"
    10  	"github.com/castai/kvisor/pkg/logging"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestEnrichmentService(t *testing.T) {
    15  	t.Run("should enrich event", func(t *testing.T) {
    16  		r := require.New(t)
    17  		wantedPodUID := "pod-uid-1"
    18  		wantedPodName := "this-is-pod-1"
    19  
    20  		svc := NewService(logging.New(&logging.Config{}), Config{
    21  			WorkerCount: 1,
    22  			EventEnrichers: []EventEnricher{
    23  				enricherFrom(func(ctx context.Context, req *EnrichRequest) {
    24  					req.Event.PodName = wantedPodName
    25  				}, castpb.EventType_UNKNOWN),
    26  				enricherFrom(func(ctx context.Context, req *EnrichRequest) {
    27  					req.Event.PodUid = wantedPodUID
    28  				}, castpb.EventType_UNKNOWN),
    29  			},
    30  		})
    31  
    32  		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    33  		defer cancel()
    34  
    35  		done := make(chan struct{})
    36  
    37  		go func() {
    38  			svc.Run(ctx)
    39  			done <- struct{}{}
    40  		}()
    41  
    42  		event := &castpb.Event{
    43  			EventType: castpb.EventType_UNKNOWN,
    44  			PodUid:    "uid-1",
    45  			PodName:   "pod-1",
    46  		}
    47  		enqueued := svc.Enqueue(&EnrichRequest{
    48  			Event:     event,
    49  			EbpfEvent: &types.Event{},
    50  		})
    51  		r.True(enqueued)
    52  
    53  		select {
    54  		case event := <-svc.Events():
    55  			r.Equal(wantedPodUID, event.PodUid)
    56  			r.Equal(wantedPodName, event.PodName)
    57  
    58  		case <-ctx.Done():
    59  			r.FailNow("timed out waiting for event")
    60  		}
    61  
    62  		cancel()
    63  
    64  		select {
    65  		case <-done:
    66  		case <-time.After(2 * time.Second):
    67  			r.FailNow("timed out waiting for service to stop")
    68  		}
    69  	})
    70  }
    71  
    72  func enricherFrom(f func(context.Context, *EnrichRequest), eventTypes ...castpb.EventType) EventEnricher {
    73  	return &functionEnricher{
    74  		f:          f,
    75  		eventTypes: eventTypes,
    76  	}
    77  }
    78  
    79  type functionEnricher struct {
    80  	f          func(context.Context, *EnrichRequest)
    81  	eventTypes []castpb.EventType
    82  }
    83  
    84  func (e *functionEnricher) Enrich(ctx context.Context, req *EnrichRequest) {
    85  	e.f(ctx, req)
    86  }
    87  
    88  func (e *functionEnricher) EventTypes() []castpb.EventType {
    89  	return e.eventTypes
    90  }