tests: increase unit test coverage for main package (#2559)

Signed-off-by: Athang69 <athangkali21@gmail.com>

Author: Athang69

KubeArmor/main.go

@@ -20,40 +20,46 @@ func init() {
 	buildinfo.PrintBuildDetails()
 }
 
-func main() {
-	if os.Geteuid() != 0 {
-		if os.Getenv("KUBEARMOR_UBI") == "" {
-			kg.Printf("Need to have root privileges to run %s\n", os.Args[0])
-			return
-		}
-	}
-	// initial clean up
-
-	bpfMapsDir := "/sys/fs/bpf/"
+func isKubeArmorBpfMap(name string) bool {
+	return strings.HasPrefix(name, "kubearmor")
+}
 
+func cleanupBpfMaps(bpfMapsDir string, removeFn func(string) error) error {
 	entries, err := os.ReadDir(bpfMapsDir)
 	if err != nil {
-		kg.Errf("Failed to read BPF map directory: %v", err)
-		return
+		return err
 	}
-
 	for _, entry := range entries {
 		if entry.IsDir() {
 			continue
 		}
-		if strings.HasPrefix(entry.Name(), "kubearmor") {
-			/* This should not be triggered in ideal cases,
-			if this is triggered that means there is incomplete cleanup process
-			from the last installation */
+		if isKubeArmorBpfMap(entry.Name()) {
 			path := filepath.Join(bpfMapsDir, entry.Name())
-			err := kl.RemoveSafe(path)
-			if err != nil {
+			if err := removeFn(path); err != nil {
 				kg.Errf("Failed to delete BPF map %s: %v", path, err)
 			} else {
 				kg.Warnf("Deleting existing map %s. This indicates previous cleanup failed", path)
 			}
 		}
 	}
+	return nil
+}
+
+func main() {
+	if os.Geteuid() != 0 {
+		if os.Getenv("KUBEARMOR_UBI") == "" {
+			kg.Printf("Need to have root privileges to run %s\n", os.Args[0])
+			return
+		}
+	}
+	// initial clean up
+
+	bpfMapsDir := "/sys/fs/bpf/"
+
+	if err := cleanupBpfMaps(bpfMapsDir, kl.RemoveSafe); err != nil {
+		kg.Errf("Failed to read BPF map directory: %v", err)
+		return
+	}
 
 	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
 	if err != nil {

KubeArmor/main_test.go

@@ -7,6 +7,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"path/filepath"
 	"strconv"
 	"testing"
 )
@@ -76,3 +77,84 @@ func TestMain(t *testing.T) {
 	main()
 	t.Log("[INFO] Terminated KubeArmor")
 }
+
+func TestBpfMapFiltering(t *testing.T) {
+	cases := []struct {
+		filename string
+		want     bool
+	}{
+		{"kubearmor_events", true},
+		{"kubearmor_policy", true},
+		{"cilium_events", false},
+		{"kube", false},
+		{"", false},
+		{"KubeArmor", false},
+	}
+	for _, tc := range cases {
+		got := isKubeArmorBpfMap(tc.filename)
+		if got != tc.want {
+			t.Errorf("isKubeArmorBpfMap(%q) = %v, want %v", tc.filename, got, tc.want)
+		}
+	}
+}
+
+func TestBpfDirCleanupWithTempDir(t *testing.T) {
+	dir := t.TempDir()
+	files := map[string]bool{
+		"kubearmor_events": true,
+		"kubearmor_policy": true,
+		"cilium_events":    false,
+	}
+	for name := range files {
+		if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), 0600); err != nil {
+			t.Fatalf("failed to create fixture file %q: %v", name, err)
+		}
+	}
+	err := cleanupBpfMaps(dir, os.Remove)
+	if err != nil {
+		t.Fatalf("cleanupBpfMaps returned unexpected error: %v", err)
+	}
+	for name, shouldDelete := range files {
+		_, err := os.Stat(filepath.Join(dir, name))
+		if shouldDelete && err == nil {
+			t.Errorf("%s should have been deleted", name)
+		}
+		if !shouldDelete && err != nil {
+			t.Errorf("%s should still exist", name)
+		}
+	}
+}
+
+func TestBpfDirCleanupSkipsDirectories(t *testing.T) {
+	dir := t.TempDir()
+	sub := filepath.Join(dir, "kubearmor_subdir")
+	if err := os.Mkdir(sub, 0755); err != nil {
+		t.Fatalf("failed to create subdir: %v", err)
+	}
+	err := cleanupBpfMaps(dir, os.Remove)
+	if err != nil {
+		t.Fatalf("cleanupBpfMaps returned unexpected error: %v", err)
+	}
+	if _, err := os.Stat(sub); os.IsNotExist(err) {
+		t.Error("subdir was incorrectly removed")
+	}
+}
+
+func TestBpfDirCleanupMissingDir(t *testing.T) {
+	dir := t.TempDir()
+	if err := os.RemoveAll(dir); err != nil {
+		t.Fatalf("failed to remove temp dir: %v", err)
+	}
+	err := cleanupBpfMaps(dir, os.Remove)
+	if err == nil {
+		t.Error("expected error for missing dir")
+	}
+}
+
+func TestNonRootWithoutUBI(t *testing.T) {
+	if os.Geteuid() == 0 {
+		t.Skip("must run as non-root")
+	}
+	t.Setenv("KUBEARMOR_UBI", "")
+	main()
+}