pexconn_test.go

   1 package torrent
   2
   3 import (
   4         "net"
   5         "testing"
   6
   7         "github.com/anacrolix/dht/v2/krpc"
   8         "github.com/stretchr/testify/require"
   9
  10         "github.com/anacrolix/torrent/metainfo"
  11         pp "github.com/anacrolix/torrent/peer_protocol"
  12 )
  13
  14 func TestPexConnState(t *testing.T) {
  15         var cl Client
  16         cl.init(TestingConfig(t))
  17         cl.initLogger()
  18         torrent := cl.newTorrent(metainfo.Hash{}, nil)
  19         addr := &net.TCPAddr{IP: net.IPv6loopback, Port: 4747}
  20         c := cl.newConnection(nil, false, addr, addr.Network(), "")
  21         c.PeerExtensionIDs = make(map[pp.ExtensionName]pp.ExtensionNumber)
  22         c.PeerExtensionIDs[pp.ExtensionNamePex] = pexExtendedId
  23         c.messageWriter.mu.Lock()
  24         c.setTorrent(torrent)
  25         if err := torrent.addPeerConn(c); err != nil {
  26                 t.Log(err)
  27         }
  28
  29         c.pex.Init(c)
  30         require.True(t, c.pex.IsEnabled(), "should get enabled")
  31         defer c.pex.Close()
  32
  33         var out pp.Message
  34         writerCalled := false
  35         testWriter := func(m pp.Message) bool {
  36                 writerCalled = true
  37                 out = m
  38                 return true
  39         }
  40         <-c.messageWriter.writeCond.Signaled()
  41         c.pex.Share(testWriter)
  42         require.True(t, writerCalled)
  43         require.EqualValues(t, pp.Extended, out.Type)
  44         require.EqualValues(t, pexExtendedId, out.ExtendedID)
  45
  46         x, err := pp.LoadPexMsg(out.ExtendedPayload)
  47         require.NoError(t, err)
  48         targx := pp.PexMsg{
  49                 Added:      krpc.CompactIPv4NodeAddrs(nil),
  50                 AddedFlags: []pp.PexPeerFlags{},
  51                 Added6: krpc.CompactIPv6NodeAddrs{
  52                         mustNodeAddr(addr),
  53                 },
  54                 Added6Flags: []pp.PexPeerFlags{0},
  55         }
  56         require.EqualValues(t, targx, x)
  57 }