From 0d77ee7a6cbdc84eb819863d0c082b3ea0d0cddd Mon Sep 17 00:00:00 2001
From: yashsingh74 <yashsingh1774@gmail.com>
Date: Wed, 29 Apr 2026 13:11:57 +0530
Subject: [PATCH] nftables: add unit tests for CleanupLeftovers

Signed-off-by: yashsingh74 <yashsingh1774@gmail.com>
---
 pkg/proxy/nftables/cleanup.go      | 28 +++++-----
 pkg/proxy/nftables/cleanup_test.go | 85 ++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 12 deletions(-)
 create mode 100644 pkg/proxy/nftables/cleanup_test.go

diff --git a/pkg/proxy/nftables/cleanup.go b/pkg/proxy/nftables/cleanup.go
index 8de243f1e45..bce9bfc78a6 100644
--- a/pkg/proxy/nftables/cleanup.go
+++ b/pkg/proxy/nftables/cleanup.go
@@ -27,23 +27,27 @@ import (
 
 // CleanupLeftovers removes all nftables rules and chains created by the Proxier
 // It returns true if an error was encountered. Errors are logged.
-func CleanupLeftovers(ctx context.Context) bool {
-	logger := klog.FromContext(ctx)
-	var encounteredError bool
-
+func CleanupLeftovers(ctx context.Context) (encounteredError bool) {
 	for _, family := range []knftables.Family{knftables.IPv4Family, knftables.IPv6Family} {
 		nft, err := knftables.New(family, kubeProxyTable)
 		if err != nil {
 			continue
 		}
-		tx := nft.NewTransaction()
-		tx.Delete(&knftables.Table{})
-		err = nft.Run(ctx, tx)
-		if err != nil && !knftables.IsNotFound(err) {
-			logger.Error(err, "Error cleaning up nftables rules")
-			encounteredError = true
-		}
+		encounteredError = cleanupLeftoversForFamily(ctx, nft) || encounteredError
+	}
+	return
+}
+
+func cleanupLeftoversForFamily(ctx context.Context, nft knftables.Interface) (encounteredError bool) {
+	logger := klog.FromContext(ctx)
+
+	tx := nft.NewTransaction()
+	tx.Delete(&knftables.Table{})
+	err := nft.Run(ctx, tx)
+	if err != nil && !knftables.IsNotFound(err) {
+		logger.Error(err, "Error cleaning up nftables rules")
+		return true
 	}
 
-	return encounteredError
+	return false
 }
diff --git a/pkg/proxy/nftables/cleanup_test.go b/pkg/proxy/nftables/cleanup_test.go
new file mode 100644
index 00000000000..b8352a48c2e
--- /dev/null
+++ b/pkg/proxy/nftables/cleanup_test.go
@@ -0,0 +1,85 @@
+//go:build linux
+
+/*
+Copyright The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package nftables
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"sigs.k8s.io/knftables"
+)
+
+func addKubeProxyTable(t *testing.T, nft *knftables.Fake) {
+	t.Helper()
+	tx := nft.NewTransaction()
+	tx.Add(&knftables.Table{})
+	if err := nft.Run(context.Background(), tx); err != nil {
+		t.Fatalf("Run: %v", err)
+	}
+}
+
+func TestCleanupLeftoversDeletesExistingTables(t *testing.T) {
+	ctx := context.Background()
+	ipv4 := knftables.NewFake(knftables.IPv4Family, kubeProxyTable)
+	ipv6 := knftables.NewFake(knftables.IPv6Family, kubeProxyTable)
+	addKubeProxyTable(t, ipv4)
+	addKubeProxyTable(t, ipv6)
+
+	if got := cleanupLeftoversForFamily(ctx, ipv4); got {
+		t.Fatalf("cleanupLeftoversForFamily(ipv4) = %v, want false", got)
+	}
+	if got := cleanupLeftoversForFamily(ctx, ipv6); got {
+		t.Fatalf("cleanupLeftoversForFamily(ipv6) = %v, want false", got)
+	}
+	if ipv4.Table != nil {
+		t.Error("IPv4 kube-proxy table still present after cleanup")
+	}
+	if ipv6.Table != nil {
+		t.Error("IPv6 kube-proxy table still present after cleanup")
+	}
+}
+
+func TestCleanupLeftoversNotFoundIgnored(t *testing.T) {
+	ctx := context.Background()
+	// No table added: delete is a no-op / NotFound on fake.
+	nft := knftables.NewFake(knftables.IPv4Family, kubeProxyTable)
+	if got := cleanupLeftoversForFamily(ctx, nft); got {
+		t.Fatalf("cleanupLeftoversForFamily() = %v, want false when table is absent", got)
+	}
+}
+
+// fakeWithRunErr embeds knftables.Fake and returns runErr from Run for testing non-NotFound failures.
+type fakeWithRunErr struct {
+	*knftables.Fake
+	runErr error
+}
+
+func (f *fakeWithRunErr) Run(ctx context.Context, tx *knftables.Transaction) error {
+	return f.runErr
+}
+
+func TestCleanupLeftoversRunErrorSetsEncountered(t *testing.T) {
+	ctx := context.Background()
+	runErr := errors.New("nft run failed")
+	nft := &fakeWithRunErr{Fake: knftables.NewFake(knftables.IPv4Family, kubeProxyTable), runErr: runErr}
+	if got := cleanupLeftoversForFamily(ctx, nft); !got {
+		t.Fatal("cleanupLeftoversForFamily() = false, want true when Run returns a non-NotFound error")
+	}
+}