github.com/Azure/aad-pod-identity@v1.8.17/pkg/cloudprovider/truncate_test.go (about)

     1  package cloudprovider
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  
     7  	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestTruncateVMIdentities(t *testing.T) {
    12  	ids := map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue{
    13  		"id-1": {},
    14  		"id-2": {},
    15  	}
    16  	truncated, rest := truncateVMIdentities(ids)
    17  	assert.Equal(t, truncated, ids)
    18  	assert.Equal(t, rest, make(map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue))
    19  
    20  	ids = make(map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue)
    21  	for i := 0; i < maxIdentitiesCount+50; i++ {
    22  		ids[fmt.Sprintf("id-%d", i)] = &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{}
    23  	}
    24  	truncated, rest = truncateVMIdentities(ids)
    25  	assert.Len(t, truncated, maxIdentitiesCount)
    26  	assert.Len(t, rest, 50)
    27  
    28  	// ensuring a key does not exist in both maps
    29  	for i := 0; i < maxIdentitiesCount+50; i++ {
    30  		key := fmt.Sprintf("id-%d", i)
    31  		_, ok1 := ids[key]
    32  		_, ok2 := rest[key]
    33  		assert.False(t, ok1 && ok2, "%s exists in both ids and rest", key)
    34  		assert.True(t, ok1 || ok2, "%s does not exist in both ids and rest", key)
    35  	}
    36  }
    37  
    38  func TestTruncateVMSSIdentities(t *testing.T) {
    39  	ids := map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue{
    40  		"id-1": {},
    41  		"id-2": {},
    42  	}
    43  	truncated, rest := truncateVMSSIdentities(ids)
    44  	assert.Equal(t, truncated, ids)
    45  	assert.Equal(t, rest, make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue))
    46  
    47  	ids = make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue)
    48  	for i := 0; i < maxIdentitiesCount+50; i++ {
    49  		ids[fmt.Sprintf("id-%d", i)] = &compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue{}
    50  	}
    51  	truncated, rest = truncateVMSSIdentities(ids)
    52  	assert.Len(t, truncated, maxIdentitiesCount)
    53  	assert.Len(t, rest, 50)
    54  
    55  	// ensuring a key does not exist in both maps
    56  	for i := 0; i < maxIdentitiesCount+50; i++ {
    57  		key := fmt.Sprintf("id-%d", i)
    58  		_, ok1 := ids[key]
    59  		_, ok2 := rest[key]
    60  		assert.False(t, ok1 && ok2, "%s exists in both ids and rest", key)
    61  		assert.True(t, ok1 || ok2, "%s does not exist in both ids and rest", key)
    62  	}
    63  }