1  package peer
  2  
  3  import (
  4  	"sync"
  5  	"testing"
  6  	"time"
  7  
  8  	"github.com/lightningnetwork/lnd/lnwire"
  9  	"github.com/stretchr/testify/require"
 10  )
 11  
 12  // TestPingManager tests three main properties about the ping manager. It
 13  // ensures that if the pong response exceeds the timeout, that a failure is
 14  // emitted on the failure channel. It ensures that if the Pong response is
 15  // not congruent with the outstanding ping then a failure is emitted on the
 16  // failure channel, and otherwise the failure channel remains empty.
 17  func TestPingManager(t *testing.T) {
 18  	t.Parallel()
 19  
 20  	testCases := []struct {
 21  		name     string
 22  		delay    int
 23  		pongSize uint16
 24  		result   bool
 25  	}{
 26  		{
 27  			name:     "happy Path",
 28  			delay:    0,
 29  			pongSize: 4,
 30  			result:   true,
 31  		},
 32  		{
 33  			name:     "bad Pong",
 34  			delay:    0,
 35  			pongSize: 3,
 36  			result:   false,
 37  		},
 38  		{
 39  			name:     "timeout",
 40  			delay:    2,
 41  			pongSize: 4,
 42  			result:   false,
 43  		},
 44  	}
 45  
 46  	payload := make([]byte, 4)
 47  	for _, test := range testCases {
 48  		t.Run(test.name, func(t *testing.T) {
 49  			// Set up PingManager.
 50  			var pingOnce sync.Once
 51  			pingSent := make(chan struct{})
 52  			disconnected := make(chan struct{})
 53  			mgr := NewPingManager(&PingManagerConfig{
 54  				NewPingPayload: func() []byte {
 55  					return payload
 56  				},
 57  				NewPongSize: func() uint16 {
 58  					return 4
 59  				},
 60  				IntervalDuration: time.Second * 2,
 61  				TimeoutDuration:  time.Second,
 62  				SendPing: func(ping *lnwire.Ping) {
 63  					pingOnce.Do(func() {
 64  						close(pingSent)
 65  					})
 66  				},
 67  				OnPongFailure: func(err error,
 68  					_ time.Duration, _ time.Duration) {
 69  
 70  					close(disconnected)
 71  				},
 72  			})
 73  			require.NoError(
 74  				t, mgr.Start(), "Could not start pingManager",
 75  			)
 76  
 77  			// Wait for initial Ping.
 78  			<-pingSent
 79  
 80  			// Wait for pre-determined time before sending Pong
 81  			// response.
 82  			time.Sleep(time.Duration(test.delay) * time.Second)
 83  
 84  			// Send Pong back.
 85  			res := lnwire.Pong{
 86  				PongBytes: make([]byte, test.pongSize),
 87  			}
 88  			mgr.ReceivedPong(&res)
 89  
 90  			select {
 91  			case <-time.NewTimer(time.Second / 2).C:
 92  				require.True(t, test.result)
 93  			case <-disconnected:
 94  				require.False(t, test.result)
 95  			}
 96  
 97  			mgr.Stop()
 98  		})
 99  	}
100  }