test: update tests

Author: santhoshmprabhu

bpf-prog/azure-block-iptables/cmd/azure-block-iptables/main.go

@@ -10,7 +10,6 @@ import (
 	"os"
 
 	"github.com/Azure/azure-container-networking/bpf-prog/azure-block-iptables/pkg/bpfprogram"
-	"github.com/cilium/ebpf/rlimit"
 )
 
 // ProgramVersion is set during build
@@ -60,11 +59,6 @@ func parseArgs() (*Config, error) {
 func attachMode(config *Config) error {
 	log.Println("Starting attach mode...")
 
-	// Remove memory limit for eBPF
-	if err := rlimit.RemoveMemlock(); err != nil {
-		return fmt.Errorf("failed to remove memlock rlimit: %w", err)
-	}
-
 	// Initialize BPF program attacher using the factory
 	bp := config.AttacherFactory()
 

bpf-prog/azure-block-iptables/cmd/azure-block-iptables/main_test.go

@@ -4,57 +4,31 @@
 package main
 
 import (
-	"os"
 	"testing"
 
 	"github.com/Azure/azure-container-networking/bpf-prog/azure-block-iptables/pkg/bpfprogram"
-	"github.com/fsnotify/fsnotify"
-	"github.com/pkg/errors"
 )
 
 func TestHandleFileEventWithMock(t *testing.T) {
 	// Create a mock Attacher
 	mockAttacher := bpfprogram.NewMockProgram()
 
-	// Create a temporary config file for testing
-	configFile := "/tmp/test-iptables-allow-list"
-
 	// Test cases
 	testCases := []struct {
 		name           string
-		setupFile      func(string) error
+		mode           string
 		expectedAttach int
 		expectedDetach int
 	}{
 		{
-			name: "empty file triggers attach",
-			setupFile: func(path string) error {
-				// Create empty file
-				file, err := os.Create(path)
-				if err != nil {
-					return errors.Wrap(err, "failed to create file")
-				}
-				return file.Close()
-			},
+			name:           "test attach mode",
+			mode:           "attach",
 			expectedAttach: 1,
 			expectedDetach: 0,
 		},
 		{
-			name: "file with content triggers detach",
-			setupFile: func(path string) error {
-				// Create file with content
-				return os.WriteFile(path, []byte("some content"), 0o600)
-			},
-			expectedAttach: 0,
-			expectedDetach: 1,
-		},
-		{
-			name: "missing file triggers detach",
-			setupFile: func(path string) error {
-				// Remove file if it exists
-				os.Remove(path)
-				return nil
-			},
+			name:           "test detach mode",
+			mode:           "detach",
 			expectedAttach: 0,
 			expectedDetach: 1,
 		},
@@ -65,20 +39,7 @@ func TestHandleFileEventWithMock(t *testing.T) {
 			// Reset mock state
 			mockAttacher.Reset()
 
-			// Setup file state
-			if err := tc.setupFile(configFile); err != nil {
-				t.Fatalf("Failed to setup file: %v", err)
-			}
-			defer os.Remove(configFile)
-
-			// Create a fake fsnotify event
-			event := fsnotify.Event{
-				Name: configFile,
-				Op:   fsnotify.Write,
-			}
-
-			// Call the function under test
-			handleFileEvent(event, configFile, mockAttacher)
+			run(&Config{Mode: tc.mode, AttacherFactory: func() bpfprogram.Attacher { return mockAttacher }})
 
 			// Verify expectations
 			if mockAttacher.AttachCallCount() != tc.expectedAttach {

bpf-prog/azure-block-iptables/pkg/bpfprogram/program.go

@@ -12,6 +12,7 @@ import (
 	blockservice "github.com/Azure/azure-container-networking/bpf-prog/azure-block-iptables/pkg/blockservice"
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/link"
+	"github.com/cilium/ebpf/rlimit"
 	"github.com/pkg/errors"
 )
 
@@ -124,6 +125,11 @@ func (p *Program) Attach() error {
 		return nil
 	}
 
+	// Remove memory limit for eBPF
+	if err := rlimit.RemoveMemlock(); err != nil {
+		return errors.Wrapf(err, "failed to remove memlock rlimit")
+	}
+
 	log.Println("Attaching BPF program...")
 
 	// Get the host network namespace inode