Add unit test for cleanup and tagging

aws · Sep 17, 2020 · 3c8e804 · 3c8e804
1 parent c6b3bd8
commit 3c8e804
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git 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
@@ -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)
+}