github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/client/pluginmanager/csimanager/usage_tracker_test.go (about)

     1  package csimanager
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/hashicorp/nomad/nomad/mock"
     7  	"github.com/hashicorp/nomad/nomad/structs"
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  func TestUsageTracker(t *testing.T) {
    12  	mockAllocs := []*structs.Allocation{
    13  		mock.Alloc(),
    14  		mock.Alloc(),
    15  		mock.Alloc(),
    16  		mock.Alloc(),
    17  		mock.Alloc(),
    18  	}
    19  
    20  	cases := []struct {
    21  		Name string
    22  
    23  		RegisterAllocs []*structs.Allocation
    24  		FreeAllocs     []*structs.Allocation
    25  
    26  		ExpectedResult bool
    27  	}{
    28  		{
    29  			Name:           "Register and deregister all allocs",
    30  			RegisterAllocs: mockAllocs,
    31  			FreeAllocs:     mockAllocs,
    32  			ExpectedResult: true,
    33  		},
    34  		{
    35  			Name:           "Register all and deregister partial allocs",
    36  			RegisterAllocs: mockAllocs,
    37  			FreeAllocs:     mockAllocs[0:3],
    38  			ExpectedResult: false,
    39  		},
    40  	}
    41  
    42  	for _, tc := range cases {
    43  		t.Run(tc.Name, func(t *testing.T) {
    44  			tracker := newVolumeUsageTracker()
    45  
    46  			volume := &structs.CSIVolume{
    47  				ID: "foo",
    48  			}
    49  			for _, alloc := range tc.RegisterAllocs {
    50  				tracker.Claim(alloc.ID, volume.ID, &UsageOptions{})
    51  			}
    52  
    53  			result := false
    54  
    55  			for _, alloc := range tc.FreeAllocs {
    56  				result = tracker.Free(alloc.ID, volume.ID, &UsageOptions{})
    57  			}
    58  
    59  			require.Equal(t, tc.ExpectedResult, result, "Tracker State: %#v", tracker.state)
    60  		})
    61  	}
    62  }