Added the enricher loop to the test code

microsoft · Dec 23, 2024 · 2703bff · 2703bff
1 parent 99b0307
commit 2703bff
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 22 deletions.
diff --git a/pkg/plugin/ebpfwindows/ebpf_windows.go b/pkg/plugin/ebpfwindows/ebpf_windows.go
@@ -34,7 +34,7 @@ var (
 type Plugin struct {
     l               *log.ZapLogger
     cfg             *kcfg.Config
-    enricher        enricher.EnricherInterface
+    enricher        *enricher.Enricher
     externalChannel chan *v1.Event
     parser          *hp.Parser
 }
@@ -240,30 +240,26 @@ func (p *Plugin) handleTraceEvent(data unsafe.Pointer, size uint64) error {
         return ErrInvalidEventData
     }
 
-    if fl != nil {
-
-        ev := &v1.Event{
-            Event:     fl,
-            Timestamp: fl.GetTime(),
-        }
+    ev := &v1.Event{
+        Event:     fl,
+        Timestamp: fl.GetTime(),
+    }
 
-        if p.enricher != nil {
-            p.enricher.Write(ev)
-        } else {
-            p.l.Error("enricher is nil when writing event")
-        }
+    if p.enricher != nil {
+        p.enricher.Write(ev)
+    } else {
+        p.l.Error("enricher is nil when writing event")
+    }
 
-        // Write the event to the external channel.
-        if p.externalChannel != nil {
-            select {
-            case p.externalChannel <- ev:
-            default:
-                // Channel is full, drop the event.
-                // We shouldn't slow down the reader.
-                metrics.LostEventsCounter.WithLabelValues(utils.ExternalChannel, name).Inc()
-            }
+    // Write the event to the external channel.
+    if p.externalChannel != nil {
+        select {
+        case p.externalChannel <- ev:
+        default:
+            // Channel is full, drop the event.
+            // We shouldn't slow down the reader.
+            metrics.LostEventsCounter.WithLabelValues(utils.ExternalChannel, name).Inc()
         }
-
     }
 
     return nil

diff --git a/pkg/plugin/ebpfwindows/ebpf_windows_test.go b/pkg/plugin/ebpfwindows/ebpf_windows_test.go
@@ -10,7 +10,10 @@ import (
     "time"
 
     kcfg "github.com/microsoft/retina/pkg/config"
+    "github.com/microsoft/retina/pkg/controllers/cache"
+    "github.com/microsoft/retina/pkg/enricher"
     "github.com/microsoft/retina/pkg/log"
+    "github.com/microsoft/retina/pkg/pubsub"
     "go.uber.org/zap"
 )
 
@@ -25,6 +28,11 @@ func TestPlugin(t *testing.T) {
         EnablePodLevel:  true,
     }
 
+    c := cache.New(pubsub.New())
+    e := enricher.New(ctx, c)
+    e.Run()
+    defer e.Reader.Close()
+
     tt := New(cfg)
 
     err := tt.Stop()