From 0391ee39339e2f4d92b9ee28d6343f6e1a5724b4 Mon Sep 17 00:00:00 2001
From: Jiri Olsa <jolsa@kernel.org>
Date: Wed, 2 Oct 2024 12:41:52 +0000
Subject: [PATCH] tetragon: Add tests for proper cleanup after sensor unload

Adding tests for proper cleanup after sensor unload.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 pkg/sensors/test/sensors_test.go | 99 ++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/pkg/sensors/test/sensors_test.go b/pkg/sensors/test/sensors_test.go
index 4d6a7e5eeaa..0b93ee33476 100644
--- a/pkg/sensors/test/sensors_test.go
+++ b/pkg/sensors/test/sensors_test.go
@@ -4,6 +4,7 @@
 package test
 
 import (
+	"os"
 	"path/filepath"
 	"testing"
 
@@ -500,3 +501,101 @@ func TestMaxEntriesInnerMulti(t *testing.T) {
 	// TODO, we need to check BTF for inner map max entries
 	t.Skip()
 }
+
+func TestCleanup(t *testing.T) {
+	p1 := program.Builder(
+		"bpf_map_test_p1.o",
+		"wake_up_new_task",
+		"kprobe/wake_up_new_task",
+		"p1",
+		"kprobe",
+	)
+	p2 := program.Builder(
+		"bpf_map_test_p2.o",
+		"wake_up_new_task",
+		"kprobe/wake_up_new_task",
+		"p2",
+		"kprobe",
+	)
+
+	var err error
+
+	verifyRemoved := func(files ...string) {
+		for _, f := range files {
+			_, err := os.Stat(filepath.Join(bpf.MapPrefixPath(), f))
+			assert.Error(t, err)
+		}
+	}
+
+	t.Run("single_ok", func(t *testing.T) {
+		m1 := program.MapBuilder("m1", p1)
+		m2 := program.MapBuilder("m2", p2)
+
+		s1 := &sensors.Sensor{
+			Name:   "sensor1",
+			Progs:  []*program.Program{p1},
+			Maps:   []*program.Map{m1, m2},
+			Policy: "policy",
+		}
+
+		err = s1.Load(bpf.MapPrefixPath())
+		assert.NoError(t, err)
+
+		s1.Unload()
+
+		verifyRemoved("m1", "m2", "policy")
+	})
+
+	t.Run("multi_ok", func(t *testing.T) {
+		m1 := program.MapBuilder("m1", p1)
+		m2 := program.MapBuilder("m2", p2)
+
+		s1 := &sensors.Sensor{
+			Name:   "sensor1",
+			Progs:  []*program.Program{p1},
+			Maps:   []*program.Map{m1, m2},
+			Policy: "policy",
+		}
+
+		s2 := &sensors.Sensor{
+			Name:   "sensor2",
+			Progs:  []*program.Program{p2},
+			Maps:   []*program.Map{m1, m2},
+			Policy: "policy",
+		}
+
+		err = s1.Load(bpf.MapPrefixPath())
+		assert.NoError(t, err)
+		err = s2.Load(bpf.MapPrefixPath())
+		assert.NoError(t, err)
+
+		s1.Unload()
+
+		verifyRemoved("policy/sensor1")
+
+		s2.Unload()
+
+		verifyRemoved("m1", "m2", "policy")
+	})
+
+	t.Run("map_fail", func(t *testing.T) {
+		m1 := program.MapBuilder("m1", p1, p2)
+		// map with wrong name
+		m2 := program.MapBuilder("non-existing", p1, p2)
+
+		s1 := &sensors.Sensor{
+			Name:   "sensor1",
+			Progs:  []*program.Program{p1},
+			Maps:   []*program.Map{m1, m2},
+			Policy: "policy",
+		}
+
+		err = s1.Load(bpf.MapPrefixPath())
+		assert.Error(t, err)
+		if err == nil {
+			defer s1.Unload()
+		}
+
+		verifyRemoved("m1", "m2", "policy")
+	})
+}