github.com/10XDev/rclone@v1.52.3-0.20200626220027-16af9ab76b2a/fs/accounting/token_bucket_test.go (about)

     1  package accounting
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/rclone/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  		"rate":           "1M",
    26  	}, out)
    27  	assert.Equal(t, rate.Limit(1048576), tokenBucket.Limit())
    28  
    29  	// Query
    30  	in = rc.Params{}
    31  	out, err = call.Fn(context.Background(), in)
    32  	require.NoError(t, err)
    33  	assert.Equal(t, rc.Params{
    34  		"bytesPerSecond": int64(1048576),
    35  		"rate":           "1M",
    36  	}, out)
    37  
    38  	// Reset
    39  	in = rc.Params{
    40  		"rate": "off",
    41  	}
    42  	out, err = call.Fn(context.Background(), in)
    43  	require.NoError(t, err)
    44  	assert.Equal(t, rc.Params{
    45  		"bytesPerSecond": int64(-1),
    46  		"rate":           "off",
    47  	}, out)
    48  	assert.Nil(t, tokenBucket)
    49  
    50  	// Query
    51  	in = rc.Params{}
    52  	out, err = call.Fn(context.Background(), in)
    53  	require.NoError(t, err)
    54  	assert.Equal(t, rc.Params{
    55  		"bytesPerSecond": int64(-1),
    56  		"rate":           "off",
    57  	}, out)
    58  
    59  }