Add unit tests for pkg/agent/util/ipset

Signed-off-by: Hongliang Liu <[email protected]>

pkg/agent/util/ipset/ipset.go

@@ -16,9 +16,10 @@ package ipset
 
 import (
 	"fmt"
-	"os/exec"
 	"regexp"
 	"strings"
+
+	"k8s.io/utils/exec"
 )
 
 type SetType string
@@ -46,16 +47,20 @@ type Interface interface {
 	ListEntries(name string) ([]string, error)
 }
 
-type Client struct{}
+type Client struct {
+	exec exec.Interface
+}
 
 var _ Interface = &Client{}
 
 func NewClient() *Client {
-	return &Client{}
+	return &Client{
+		exec: exec.New(),
+	}
 }
 
 func (c *Client) DestroyIPSet(name string) error {
-	cmd := exec.Command("ipset", "destroy", name)
+	cmd := c.exec.Command("ipset", "destroy", name)
 	if output, err := cmd.CombinedOutput(); err != nil {
 		if strings.Contains(err.Error(), "The set with the given name does not exist") {
 			return nil
@@ -67,13 +72,13 @@ func (c *Client) DestroyIPSet(name string) error {
 
 // CreateIPSet creates a new set, it will ignore error when the set already exists.
 func (c *Client) CreateIPSet(name string, setType SetType, isIPv6 bool) error {
-	var cmd *exec.Cmd
+	var cmd exec.Cmd
 	if isIPv6 {
 		// #nosec G204 -- inputs are not controlled by users
-		cmd = exec.Command("ipset", "create", name, string(setType), "family", "inet6", "-exist")
+		cmd = c.exec.Command("ipset", "create", name, string(setType), "family", "inet6", "-exist")
 	} else {
 		// #nosec G204 -- inputs are not controlled by users
-		cmd = exec.Command("ipset", "create", name, string(setType), "-exist")
+		cmd = c.exec.Command("ipset", "create", name, string(setType), "-exist")
 	}
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("error creating ipset %s, err: %w, output: %s", name, err, output)
@@ -83,7 +88,7 @@ func (c *Client) CreateIPSet(name string, setType SetType, isIPv6 bool) error {
 
 // AddEntry adds a new entry to the set, it will ignore error when the entry already exists.
 func (c *Client) AddEntry(name string, entry string) error {
-	cmd := exec.Command("ipset", "add", name, entry, "-exist")
+	cmd := c.exec.Command("ipset", "add", name, entry, "-exist")
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("error adding entry %s to ipset %s, err: %w, output: %s", entry, name, err, output)
 	}
@@ -92,7 +97,7 @@ func (c *Client) AddEntry(name string, entry string) error {
 
 // DelEntry deletes the entry from the set, it will ignore error when the entry doesn't exist.
 func (c *Client) DelEntry(name string, entry string) error {
-	cmd := exec.Command("ipset", "del", name, entry, "-exist")
+	cmd := c.exec.Command("ipset", "del", name, entry, "-exist")
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("error deleting entry %s from ipset %s, err: %w, output: %s", entry, name, err, output)
 	}
@@ -101,7 +106,7 @@ func (c *Client) DelEntry(name string, entry string) error {
 
 // ListEntries lists all the entries of the set.
 func (c *Client) ListEntries(name string) ([]string, error) {
-	cmd := exec.Command("ipset", "list", name)
+	cmd := c.exec.Command("ipset", "list", name)
 	output, err := cmd.CombinedOutput()
 	if err != nil {
 		return nil, fmt.Errorf("error listing ipset %s, err: %w, output: %s", name, err, output)

pkg/agent/util/ipset/ipset_test.go

@@ -0,0 +1,219 @@
+// Copyright 2024 Antrea 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 ipset
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"k8s.io/utils/exec"
+	exectesting "k8s.io/utils/exec/testing"
+)
+
+const testIPSetName = "test-ipset"
+
+func TestClient_CreateIPSet(t *testing.T) {
+	tests := []struct {
+		name              string
+		setType           SetType
+		isIPv6            bool
+		err               error
+		output            string
+		expectedArgs      []string
+		expectedErrOutput string
+	}{
+		{
+			name:         "Create IPv4 hash:net set",
+			setType:      HashNet,
+			expectedArgs: []string{"create", testIPSetName, string(HashNet), "-exist"},
+		},
+		{
+			name:         "Create IPv6 hash:ip set",
+			setType:      HashIP,
+			isIPv6:       true,
+			expectedArgs: []string{"create", testIPSetName, string(HashIP), "family", "inet6", "-exist"},
+		},
+		{
+			name:              "Create IPv4 set with error",
+			setType:           HashIPPort,
+			err:               errors.New("some errors"),
+			output:            "error output",
+			expectedArgs:      []string{"create", testIPSetName, string(HashIPPort), "-exist"},
+			expectedErrOutput: fmt.Sprintf("error creating ipset %s, err: some errors, output: error output", testIPSetName),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			fcmd := exectesting.FakeCmd{
+				CombinedOutputScript: []exectesting.FakeAction{
+					func() ([]byte, []byte, error) { return []byte(tt.output), nil, tt.err },
+				},
+			}
+			fakeExec := &exectesting.FakeExec{
+				CommandScript: []exectesting.FakeCommandAction{
+					func(cmd string, args ...string) exec.Cmd {
+						assert.Equal(t, "ipset", cmd)
+						assert.Equal(t, tt.expectedArgs, args)
+						return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+					},
+				},
+			}
+			c := &Client{exec: fakeExec}
+			err := c.CreateIPSet(testIPSetName, tt.setType, tt.isIPv6)
+			if tt.expectedErrOutput != "" {
+				assert.EqualError(t, err, tt.expectedErrOutput)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestClient_DestroyIPSet(t *testing.T) {
+	tests := []struct {
+		name              string
+		err               error
+		output            string
+		expectedErrOutput string
+	}{
+		{
+			name: "Destroy set successfully",
+		},
+		{
+			name: "Destroy non-existent set, no error",
+			err:  errors.New("The set with the given name does not exist"),
+		},
+		{
+			name:              "Destroy set with other error",
+			err:               errors.New("some errors"),
+			output:            "error output",
+			expectedErrOutput: fmt.Sprintf("error destroying ipset %s, err: some errors, output: error output", testIPSetName),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			fcmd := exectesting.FakeCmd{
+				CombinedOutputScript: []exectesting.FakeAction{
+					func() ([]byte, []byte, error) { return []byte(tt.output), nil, tt.err },
+				},
+			}
+			fakeExec := &exectesting.FakeExec{
+				CommandScript: []exectesting.FakeCommandAction{
+					func(cmd string, args ...string) exec.Cmd {
+						assert.Equal(t, "ipset", cmd)
+						assert.Equal(t, []string{"destroy", testIPSetName}, args)
+						return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+					},
+				},
+			}
+			c := &Client{exec: fakeExec}
+			err := c.DestroyIPSet(testIPSetName)
+			if tt.expectedErrOutput != "" {
+				assert.EqualError(t, err, tt.expectedErrOutput)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestClient_AddEntry(t *testing.T) {
+	fcmd := exectesting.FakeCmd{
+		CombinedOutputScript: []exectesting.FakeAction{
+			func() ([]byte, []byte, error) { return []byte("error output"), nil, errors.New("some errors") },
+			func() ([]byte, []byte, error) { return nil, nil, nil },
+		},
+	}
+	fakeExec := &exectesting.FakeExec{
+		CommandScript: []exectesting.FakeCommandAction{
+			func(cmd string, args ...string) exec.Cmd {
+				assert.Equal(t, "ipset", cmd)
+				assert.Equal(t, []string{"add", testIPSetName, "1..2.3.4", "-exist"}, args)
+				return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+			},
+			func(cmd string, args ...string) exec.Cmd {
+				assert.Equal(t, "ipset", cmd)
+				assert.Equal(t, []string{"add", testIPSetName, "1.2.3.4", "-exist"}, args)
+				return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+			},
+		},
+	}
+	c := &Client{exec: fakeExec}
+	err := c.AddEntry(testIPSetName, "1..2.3.4")
+	assert.EqualError(t, err, fmt.Sprintf("error adding entry 1..2.3.4 to ipset %s, err: some errors, output: error output", testIPSetName))
+	err = c.AddEntry(testIPSetName, "1.2.3.4")
+	assert.NoError(t, err)
+}
+
+func TestClient_DelEntry(t *testing.T) {
+	fcmd := exectesting.FakeCmd{
+		CombinedOutputScript: []exectesting.FakeAction{
+			func() ([]byte, []byte, error) { return []byte("error output"), nil, errors.New("some errors") },
+			func() ([]byte, []byte, error) { return nil, nil, nil },
+		},
+	}
+	fakeExec := &exectesting.FakeExec{
+		CommandScript: []exectesting.FakeCommandAction{
+			func(cmd string, args ...string) exec.Cmd {
+				assert.Equal(t, "ipset", cmd)
+				assert.Equal(t, []string{"del", testIPSetName, "1..2.3.4", "-exist"}, args)
+				return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+			},
+			func(cmd string, args ...string) exec.Cmd {
+				assert.Equal(t, "ipset", cmd)
+				assert.Equal(t, []string{"del", testIPSetName, "1.2.3.4", "-exist"}, args)
+				return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+			},
+		},
+	}
+	c := &Client{exec: fakeExec}
+	err := c.DelEntry(testIPSetName, "1..2.3.4")
+	assert.EqualError(t, err, fmt.Sprintf("error deleting entry 1..2.3.4 from ipset %s, err: some errors, output: error output", testIPSetName))
+	err = c.DelEntry(testIPSetName, "1.2.3.4")
+	assert.NoError(t, err)
+}
+
+func TestClient_ListEntries(t *testing.T) {
+	fcmd := exectesting.FakeCmd{
+		CombinedOutputScript: []exectesting.FakeAction{
+			func() ([]byte, []byte, error) { return []byte("error output"), nil, errors.New("some errors") },
+			func() ([]byte, []byte, error) { return []byte("1.1.1.1\n2.2.2.2"), nil, nil },
+		},
+	}
+	fakeExec := &exectesting.FakeExec{
+		CommandScript: []exectesting.FakeCommandAction{
+			func(cmd string, args ...string) exec.Cmd {
+				assert.Equal(t, "ipset", cmd)
+				assert.Equal(t, []string{"list", testIPSetName}, args)
+				return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+			},
+			func(cmd string, args ...string) exec.Cmd {
+				assert.Equal(t, "ipset", cmd)
+				assert.Equal(t, []string{"list", testIPSetName}, args)
+				return exectesting.InitFakeCmd(&fcmd, cmd, args...)
+			},
+		},
+	}
+	c := &Client{exec: fakeExec}
+	entries, err := c.ListEntries(testIPSetName)
+	assert.EqualError(t, err, fmt.Sprintf("error listing ipset %s, err: some errors, output: error output", testIPSetName))
+	assert.Nil(t, entries)
+	entries, err = c.ListEntries(testIPSetName)
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"1.1.1.1", "2.2.2.2"}, entries)
+}