From 541fd58791111374d5a54e0cc47f1129a727fc24 Mon Sep 17 00:00:00 2001
From: Penny Zheng <penny.zheng@arm.com>
Date: Wed, 17 Jun 2020 02:55:59 +0000
Subject: [PATCH] rate-limiter: add rate limiter unit test

add TestRxRateLimiter and TestTxRateLimiter unit tests

Fixes: #250

Signed-off-by: Penny Zheng <penny.zheng@arm.com>
---
 src/runtime/virtcontainers/network_test.go | 99 ++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/src/runtime/virtcontainers/network_test.go b/src/runtime/virtcontainers/network_test.go
index b8d3a6b3f9..cbb303274f 100644
--- a/src/runtime/virtcontainers/network_test.go
+++ b/src/runtime/virtcontainers/network_test.go
@@ -12,6 +12,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/containernetworking/plugins/pkg/ns"
 	ktu "github.com/kata-containers/kata-containers/src/runtime/pkg/katatestutils"
 	vcTypes "github.com/kata-containers/kata-containers/src/runtime/virtcontainers/pkg/types"
 	"github.com/stretchr/testify/assert"
@@ -284,3 +285,101 @@ func TestTcRedirectNetwork(t *testing.T) {
 	err = netHandle.LinkDel(link)
 	assert.NoError(err)
 }
+
+func TestRxRateLimiter(t *testing.T) {
+	if tc.NotValid(ktu.NeedRoot()) {
+		t.Skip(testDisabledAsNonRoot)
+	}
+
+	assert := assert.New(t)
+
+	netHandle, err := netlink.NewHandle()
+	assert.NoError(err)
+	defer netHandle.Delete()
+
+	// Create a test veth interface.
+	vethName := "foo"
+	veth := &netlink.Veth{LinkAttrs: netlink.LinkAttrs{Name: vethName, TxQLen: 200, MTU: 1400}, PeerName: "bar"}
+
+	err = netlink.LinkAdd(veth)
+	assert.NoError(err)
+
+	endpoint, err := createVethNetworkEndpoint(1, vethName, NetXConnectTCFilterModel)
+	assert.NoError(err)
+
+	link, err := netlink.LinkByName(vethName)
+	assert.NoError(err)
+
+	err = netHandle.LinkSetUp(link)
+	assert.NoError(err)
+
+	err = setupTCFiltering(endpoint, 1, true)
+	assert.NoError(err)
+
+	// 10Mb
+	maxRate := uint64(10000000)
+	err = addRxRateLimiter(endpoint, maxRate)
+	assert.NoError(err)
+
+	currentNS, err := ns.GetCurrentNS()
+	assert.NoError(err)
+
+	err = removeRxRateLimiter(endpoint, currentNS.Path())
+	assert.NoError(err)
+
+	err = removeTCFiltering(endpoint)
+	assert.NoError(err)
+
+	// Remove the veth created for testing.
+	err = netHandle.LinkDel(link)
+	assert.NoError(err)
+}
+
+func TestTxRateLimiter(t *testing.T) {
+	if tc.NotValid(ktu.NeedRoot()) {
+		t.Skip(testDisabledAsNonRoot)
+	}
+
+	assert := assert.New(t)
+
+	netHandle, err := netlink.NewHandle()
+	assert.NoError(err)
+	defer netHandle.Delete()
+
+	// Create a test veth interface.
+	vethName := "foo"
+	veth := &netlink.Veth{LinkAttrs: netlink.LinkAttrs{Name: vethName, TxQLen: 200, MTU: 1400}, PeerName: "bar"}
+
+	err = netlink.LinkAdd(veth)
+	assert.NoError(err)
+
+	endpoint, err := createVethNetworkEndpoint(1, vethName, NetXConnectTCFilterModel)
+	assert.NoError(err)
+
+	link, err := netlink.LinkByName(vethName)
+	assert.NoError(err)
+
+	err = netHandle.LinkSetUp(link)
+	assert.NoError(err)
+
+	err = setupTCFiltering(endpoint, 1, true)
+	assert.NoError(err)
+
+	// 10Mb
+	maxRate := uint64(10000000)
+	err = addTxRateLimiter(endpoint, maxRate)
+	assert.NoError(err)
+
+	currentNS, err := ns.GetCurrentNS()
+	assert.NoError(err)
+
+	err = removeTxRateLimiter(endpoint, currentNS.Path())
+	assert.NoError(err)
+
+	err = removeTCFiltering(endpoint)
+	assert.NoError(err)
+
+	// Remove the veth created for testing.
+	err = netHandle.LinkDel(link)
+	assert.NoError(err)
+}