github.com/divyam234/rclone@v1.64.1/fs/accounting/token_bucket_test.go (about)

     1  package accounting
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/divyam234/rclone/fs/rc"
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/stretchr/testify/require"
    10  	"golang.org/x/time/rate"
    11  )
    12  
    13  func TestRcBwLimit(t *testing.T) {
    14  	call := rc.Calls.Get("core/bwlimit")
    15  	assert.NotNil(t, call)
    16  
    17  	// Set
    18  	in := rc.Params{
    19  		"rate": "1M",
    20  	}
    21  	out, err := call.Fn(context.Background(), in)
    22  	require.NoError(t, err)
    23  	assert.Equal(t, rc.Params{
    24  		"bytesPerSecond":   int64(1048576),
    25  		"bytesPerSecondTx": int64(1048576),
    26  		"bytesPerSecondRx": int64(1048576),
    27  		"rate":             "1Mi",
    28  	}, out)
    29  	assert.Equal(t, rate.Limit(1048576), TokenBucket.curr[0].Limit())
    30  
    31  	// Query
    32  	in = rc.Params{}
    33  	out, err = call.Fn(context.Background(), in)
    34  	require.NoError(t, err)
    35  	assert.Equal(t, rc.Params{
    36  		"bytesPerSecond":   int64(1048576),
    37  		"bytesPerSecondTx": int64(1048576),
    38  		"bytesPerSecondRx": int64(1048576),
    39  		"rate":             "1Mi",
    40  	}, out)
    41  
    42  	// Set
    43  	in = rc.Params{
    44  		"rate": "10M:1M",
    45  	}
    46  	out, err = call.Fn(context.Background(), in)
    47  	require.NoError(t, err)
    48  	assert.Equal(t, rc.Params{
    49  		"bytesPerSecond":   int64(10485760),
    50  		"bytesPerSecondTx": int64(10485760),
    51  		"bytesPerSecondRx": int64(1048576),
    52  		"rate":             "10Mi:1Mi",
    53  	}, out)
    54  	assert.Equal(t, rate.Limit(10485760), TokenBucket.curr[0].Limit())
    55  
    56  	// Query
    57  	in = rc.Params{}
    58  	out, err = call.Fn(context.Background(), in)
    59  	require.NoError(t, err)
    60  	assert.Equal(t, rc.Params{
    61  		"bytesPerSecond":   int64(10485760),
    62  		"bytesPerSecondTx": int64(10485760),
    63  		"bytesPerSecondRx": int64(1048576),
    64  		"rate":             "10Mi:1Mi",
    65  	}, out)
    66  
    67  	// Reset
    68  	in = rc.Params{
    69  		"rate": "off",
    70  	}
    71  	out, err = call.Fn(context.Background(), in)
    72  	require.NoError(t, err)
    73  	assert.Equal(t, rc.Params{
    74  		"bytesPerSecond":   int64(-1),
    75  		"bytesPerSecondTx": int64(-1),
    76  		"bytesPerSecondRx": int64(-1),
    77  		"rate":             "off",
    78  	}, out)
    79  	assert.Nil(t, TokenBucket.curr[0])
    80  
    81  	// Query
    82  	in = rc.Params{}
    83  	out, err = call.Fn(context.Background(), in)
    84  	require.NoError(t, err)
    85  	assert.Equal(t, rc.Params{
    86  		"bytesPerSecond":   int64(-1),
    87  		"bytesPerSecondTx": int64(-1),
    88  		"bytesPerSecondRx": int64(-1),
    89  		"rate":             "off",
    90  	}, out)
    91  
    92  }