github.com/uber/kraken@v0.1.4/tracker/peerstore/redis_test.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package peerstore
    15  
    16  import (
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/uber/kraken/core"
    21  
    22  	"github.com/alicebob/miniredis"
    23  	"github.com/andres-erbsen/clock"
    24  	"github.com/stretchr/testify/require"
    25  )
    26  
    27  func redisConfigFixture() RedisConfig {
    28  	s, err := miniredis.Run()
    29  	if err != nil {
    30  		panic(err)
    31  	}
    32  	return RedisConfig{
    33  		Addr:              s.Addr(),
    34  		PeerSetWindowSize: 30 * time.Second,
    35  		MaxPeerSetWindows: 4,
    36  	}
    37  }
    38  
    39  func TestRedisStoreGetPeersPopulatesPeerInfoFields(t *testing.T) {
    40  	require := require.New(t)
    41  
    42  	config := redisConfigFixture()
    43  
    44  	s, err := NewRedisStore(config, clock.New())
    45  	require.NoError(err)
    46  
    47  	h := core.InfoHashFixture()
    48  
    49  	p := core.PeerInfoFixture()
    50  	p.Complete = true
    51  
    52  	require.NoError(s.UpdatePeer(h, p))
    53  
    54  	peers, err := s.GetPeers(h, 1)
    55  	require.NoError(err)
    56  	require.Equal(peers, []*core.PeerInfo{p})
    57  }
    58  
    59  func TestRedisStoreGetPeersFromMultipleWindows(t *testing.T) {
    60  	require := require.New(t)
    61  
    62  	config := redisConfigFixture()
    63  	config.PeerSetWindowSize = 10 * time.Second
    64  	config.MaxPeerSetWindows = 3
    65  
    66  	clk := clock.NewMock()
    67  	clk.Set(time.Now())
    68  
    69  	s, err := NewRedisStore(config, clk)
    70  	require.NoError(err)
    71  
    72  	// Reset time to the beginning of a window.
    73  	clk.Set(time.Unix(s.curPeerSetWindow(), 0))
    74  
    75  	h := core.InfoHashFixture()
    76  
    77  	// Each peer will be added on a different second to distribute them across
    78  	// multiple windows.
    79  	var peers []*core.PeerInfo
    80  	for i := 0; i < int(config.PeerSetWindowSize.Seconds())*config.MaxPeerSetWindows; i++ {
    81  		if i > 0 {
    82  			clk.Add(time.Second)
    83  		}
    84  		p := core.PeerInfoFixture()
    85  		peers = append(peers, p)
    86  		require.NoError(s.UpdatePeer(h, p))
    87  	}
    88  
    89  	result, err := s.GetPeers(h, len(peers))
    90  	require.NoError(err)
    91  	require.Equal(core.SortedByPeerID(peers), core.SortedByPeerID(result))
    92  }
    93  
    94  func TestRedisStoreGetPeersLimit(t *testing.T) {
    95  	require := require.New(t)
    96  
    97  	config := redisConfigFixture()
    98  	config.PeerSetWindowSize = 10 * time.Second
    99  	config.MaxPeerSetWindows = 3
   100  
   101  	clk := clock.NewMock()
   102  	clk.Set(time.Now())
   103  
   104  	s, err := NewRedisStore(config, clk)
   105  	require.NoError(err)
   106  
   107  	// Reset time to the beginning of a window.
   108  	clk.Set(time.Unix(s.curPeerSetWindow(), 0))
   109  
   110  	h := core.InfoHashFixture()
   111  
   112  	// Each peer will be added on a different second to distribute them across
   113  	// multiple windows.
   114  	for i := 0; i < 30; i++ {
   115  		if i > 0 {
   116  			clk.Add(time.Second)
   117  		}
   118  		require.NoError(s.UpdatePeer(h, core.PeerInfoFixture()))
   119  	}
   120  
   121  	// Request more peers than were added on a single window to ensure we obey the limit
   122  	// across multiple windows.
   123  	for i := 0; i < 100; i++ {
   124  		result, err := s.GetPeers(h, 15)
   125  		require.NoError(err)
   126  		require.Len(result, 15)
   127  	}
   128  }
   129  
   130  func TestRedisStoreGetPeersCollapsesCompleteBits(t *testing.T) {
   131  	require := require.New(t)
   132  
   133  	config := redisConfigFixture()
   134  
   135  	s, err := NewRedisStore(config, clock.New())
   136  	require.NoError(err)
   137  
   138  	h := core.InfoHashFixture()
   139  	p := core.PeerInfoFixture()
   140  
   141  	require.NoError(s.UpdatePeer(h, p))
   142  
   143  	peers, err := s.GetPeers(h, 2)
   144  	require.NoError(err)
   145  	require.Len(peers, 1)
   146  	require.False(peers[0].Complete)
   147  
   148  	p.Complete = true
   149  	require.NoError(s.UpdatePeer(h, p))
   150  
   151  	peers, err = s.GetPeers(h, 2)
   152  	require.NoError(err)
   153  	require.Len(peers, 1)
   154  	require.True(peers[0].Complete)
   155  }
   156  
   157  func TestRedisStorePeerExpiration(t *testing.T) {
   158  	require := require.New(t)
   159  
   160  	config := redisConfigFixture()
   161  	config.PeerSetWindowSize = time.Second
   162  	config.MaxPeerSetWindows = 2
   163  
   164  	s, err := NewRedisStore(config, clock.New())
   165  	require.NoError(err)
   166  
   167  	h := core.InfoHashFixture()
   168  	p := core.PeerInfoFixture()
   169  
   170  	require.NoError(s.UpdatePeer(h, p))
   171  
   172  	result, err := s.GetPeers(h, 1)
   173  	require.NoError(err)
   174  	require.Len(result, 1)
   175  
   176  	time.Sleep(3 * time.Second)
   177  
   178  	result, err = s.GetPeers(h, 1)
   179  	require.NoError(err)
   180  	require.Empty(result)
   181  }