github.com/status-im/status-go@v1.1.0/mailserver/limiter_test.go (about)

     1  package mailserver
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestIsAllowed(t *testing.T) {
    12  	peerID := "peerID"
    13  	testCases := []struct {
    14  		t               time.Duration
    15  		shouldBeAllowed bool
    16  		db              func() map[string]time.Time
    17  		info            string
    18  	}{
    19  		{
    20  			t:               5 * time.Millisecond,
    21  			shouldBeAllowed: true,
    22  			db: func() map[string]time.Time {
    23  				return make(map[string]time.Time)
    24  			},
    25  			info: "Expecting limiter.isAllowed to allow with an empty db",
    26  		},
    27  		{
    28  			t:               5 * time.Millisecond,
    29  			shouldBeAllowed: true,
    30  			db: func() map[string]time.Time {
    31  				db := make(map[string]time.Time)
    32  				db[peerID] = time.Now().Add(time.Duration(-10) * time.Millisecond)
    33  				return db
    34  			},
    35  			info: "Expecting limiter.isAllowed to allow with an expired peer on its db",
    36  		},
    37  		{
    38  			t:               5 * time.Millisecond,
    39  			shouldBeAllowed: false,
    40  			db: func() map[string]time.Time {
    41  				db := make(map[string]time.Time)
    42  				db[peerID] = time.Now().Add(time.Duration(-1) * time.Millisecond)
    43  				return db
    44  			},
    45  			info: "Expecting limiter.isAllowed to not allow with a non expired peer on its db",
    46  		},
    47  	}
    48  
    49  	for _, tc := range testCases {
    50  		t.Run(tc.info, func(*testing.T) {
    51  			l := newRateLimiter(tc.t)
    52  			l.db = tc.db()
    53  			assert.Equal(t, tc.shouldBeAllowed, l.IsAllowed(peerID))
    54  		})
    55  	}
    56  }
    57  
    58  func TestRemoveExpiredRateLimits(t *testing.T) {
    59  	peer := "peer"
    60  	l := newRateLimiter(time.Duration(5) * time.Second)
    61  	for i := 0; i < 10; i++ {
    62  		peerID := fmt.Sprintf("%s%d", peer, i)
    63  		l.db[peerID] = time.Now().Add(time.Duration(i*(-2)) * time.Second)
    64  	}
    65  
    66  	l.deleteExpired()
    67  	assert.Equal(t, 3, len(l.db))
    68  
    69  	for i := 0; i < 3; i++ {
    70  		peerID := fmt.Sprintf("%s%d", peer, i)
    71  		_, ok := l.db[peerID]
    72  		assert.True(t, ok, fmt.Sprintf("Non expired peer '%s' should exist, but it doesn't", peerID))
    73  	}
    74  	for i := 3; i < 10; i++ {
    75  		peerID := fmt.Sprintf("%s%d", peer, i)
    76  		_, ok := l.db[peerID]
    77  		assert.False(t, ok, fmt.Sprintf("Expired peer '%s' should not exist, but it does", peerID))
    78  	}
    79  }
    80  
    81  func TestCleaningUpExpiredRateLimits(t *testing.T) {
    82  	l := newRateLimiter(5 * time.Second)
    83  	l.period = time.Millisecond * 10
    84  	l.Start()
    85  	defer l.Stop()
    86  
    87  	l.db["peer01"] = time.Now().Add(-1 * time.Second)
    88  	l.db["peer02"] = time.Now().Add(-2 * time.Second)
    89  	l.db["peer03"] = time.Now().Add(-10 * time.Second)
    90  
    91  	time.Sleep(time.Millisecond * 20)
    92  
    93  	_, ok := l.db["peer01"]
    94  	assert.True(t, ok)
    95  	_, ok = l.db["peer02"]
    96  	assert.True(t, ok)
    97  	_, ok = l.db["peer03"]
    98  	assert.False(t, ok)
    99  }
   100  
   101  func TestAddingLimts(t *testing.T) {
   102  	peerID := "peerAdding"
   103  	l := newRateLimiter(time.Duration(5) * time.Second)
   104  	pre := time.Now()
   105  	l.Add(peerID)
   106  	post := time.Now()
   107  	assert.True(t, l.db[peerID].After(pre))
   108  	assert.True(t, l.db[peerID].Before(post))
   109  }