github.com/koko1123/flow-go-1@v0.29.6/network/p2p/blob/rate_limit_test.go (about)

     1  package blob
     2  
     3  import (
     4  	"context"
     5  	"crypto/rand"
     6  	"testing"
     7  
     8  	blocks "github.com/ipfs/go-block-format"
     9  	"github.com/ipfs/go-datastore"
    10  	dssync "github.com/ipfs/go-datastore/sync"
    11  	blockstore "github.com/ipfs/go-ipfs-blockstore"
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestRateLimit(t *testing.T) {
    17  	bs := blockstore.NewBlockstore(dssync.MutexWrap(datastore.NewMapDatastore()))
    18  	rateLimitedBs := newRateLimitedBlockStore(bs, "rate_test", 4, 4) // 4 bytes per second, burst of 4
    19  
    20  	data := make([]byte, 4)
    21  	_, _ = rand.Read(data)
    22  	blk := blocks.NewBlock(data)
    23  	err := rateLimitedBs.Put(context.Background(), blk)
    24  	require.NoError(t, err)
    25  
    26  	_, err = rateLimitedBs.Get(context.Background(), blk.Cid())
    27  	assert.NoError(t, err)
    28  
    29  	_, err = rateLimitedBs.Get(context.Background(), blk.Cid())
    30  	assert.ErrorIs(t, err, rateLimitedError) // second request should be rate limited
    31  }
    32  
    33  func TestBurstLimit(t *testing.T) {
    34  	bs := blockstore.NewBlockstore(dssync.MutexWrap(datastore.NewMapDatastore()))
    35  	rateLimitedBs := newRateLimitedBlockStore(bs, "burst_test", 4, 8) // 4 bytes per second, burst of 8
    36  
    37  	data := make([]byte, 4)
    38  	_, _ = rand.Read(data)
    39  	blk := blocks.NewBlock(data)
    40  	err := rateLimitedBs.Put(context.Background(), blk)
    41  	require.NoError(t, err)
    42  
    43  	_, err = rateLimitedBs.Get(context.Background(), blk.Cid())
    44  	assert.NoError(t, err)
    45  
    46  	_, err = rateLimitedBs.Get(context.Background(), blk.Cid())
    47  	assert.NoError(t, err) // second request is allowed due to burst limit
    48  }