From 3c8e8046c2f8e31e01f5bd818980954844bc333b Mon Sep 17 00:00:00 2001
From: Claes Mogren <claes.mogren@gmail.com>
Date: Fri, 11 Sep 2020 22:16:51 -0700
Subject: [PATCH] Add unit test for cleanup and tagging

---
 pkg/awsutils/awsutils.go      |  5 ++++-
 pkg/awsutils/awsutils_test.go | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/pkg/awsutils/awsutils.go b/pkg/awsutils/awsutils.go
index 8165e786cb..43cf45ab06 100644
--- a/pkg/awsutils/awsutils.go
+++ b/pkg/awsutils/awsutils.go
@@ -1413,8 +1413,11 @@ func (cache *EC2InstanceMetadataCache) DeallocIPAddresses(eniID string, ips []st
 }
 
 func (cache *EC2InstanceMetadataCache) cleanUpLeakedENIs() {
+	cache.cleanUpLeakedENIsInternal(time.Duration(rand.Intn(eniCleanupStartupDelayMax)) * time.Second)
+}
+
+func (cache *EC2InstanceMetadataCache) cleanUpLeakedENIsInternal(startupDelay time.Duration) {
 	rand.Seed(time.Now().UnixNano())
-	startupDelay := time.Duration(rand.Intn(eniCleanupStartupDelayMax)) * time.Second
 	log.Infof("Will attempt to clean up AWS CNI leaked ENIs after waiting %s.", startupDelay)
 	time.Sleep(startupDelay)
 
diff --git a/pkg/awsutils/awsutils_test.go b/pkg/awsutils/awsutils_test.go
index 5befa68338..f5af5228ee 100644
--- a/pkg/awsutils/awsutils_test.go
+++ b/pkg/awsutils/awsutils_test.go
@@ -890,3 +890,25 @@ func TestEC2InstanceMetadataCache_SetUnmanagedENIs(t *testing.T) {
 	ins.SetUnmanagedENIs(nil)
 	assert.False(t, ins.IsUnmanagedENI("eni-1"))
 }
+
+func TestEC2InstanceMetadataCache_cleanUpLeakedENIsInternal(t *testing.T) {
+	ctrl, _, mockEC2 := setup(t)
+	defer ctrl.Finish()
+
+	description := eniDescriptionPrefix + "test"
+	result := &ec2.DescribeNetworkInterfacesOutput{
+		NetworkInterfaces: []*ec2.NetworkInterface{{
+			Description: &description,
+			TagSet: []*ec2.Tag{
+				{Key: aws.String(eniNodeTagKey), Value: aws.String("test-value")},
+			},
+		}},
+	}
+
+	mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil)
+	mockEC2.EXPECT().CreateTagsWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
+
+	ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2}
+	// Test checks that both mocks gets called.
+	ins.cleanUpLeakedENIsInternal(time.Millisecond)
+}