github.com/Bytom/bytom@v1.1.2-0.20210127130405-ae40204c0b09/p2p/discover/mdns/lan_discover_test.go (about)

     1  package mdns
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/davecgh/go-spew/spew"
     7  	"net"
     8  	"reflect"
     9  	"time"
    10  )
    11  
    12  var wantEvents = []LANPeerEvent{
    13  	{IP: []net.IP{net.IPv4(1, 2, 3, 4)}, Port: 1024},
    14  	{IP: []net.IP{net.IPv4(1, 2, 3, 4), net.IPv4(5, 6, 7, 8)}, Port: 1024},
    15  }
    16  
    17  type mockProtocol struct {
    18  	entries chan *LANPeerEvent
    19  }
    20  
    21  func newMockProtocol() *mockProtocol {
    22  	return &mockProtocol{
    23  		entries: make(chan *LANPeerEvent, 1024),
    24  	}
    25  }
    26  func (m *mockProtocol) registerService(port int) error {
    27  	return nil
    28  }
    29  
    30  func (m *mockProtocol) registerResolver(event chan LANPeerEvent) error {
    31  	for _, peerEvent := range wantEvents {
    32  		event <- peerEvent
    33  	}
    34  	return nil
    35  }
    36  
    37  func (m *mockProtocol) stopService() {
    38  
    39  }
    40  
    41  func (m *mockProtocol) stopResolver() {
    42  
    43  }
    44  
    45  func TestLanDiscover(t *testing.T) {
    46  	lanDiscv := NewLANDiscover(newMockProtocol(), 12345)
    47  	defer lanDiscv.Stop()
    48  
    49  	lanPeerMsgSub, err := lanDiscv.Subscribe()
    50  	if err != nil {
    51  		t.Fatal("subscribe lan peer msg err")
    52  	}
    53  
    54  	var gotevents = []LANPeerEvent{}
    55  	timeout := time.After(1 * time.Second)
    56  	for {
    57  		select {
    58  		case obj, ok := <-lanPeerMsgSub.Chan():
    59  			if !ok {
    60  				t.Fatal("subscription channel closed")
    61  				return
    62  			}
    63  
    64  			ev, ok := obj.Data.(LANPeerEvent)
    65  			if !ok {
    66  				t.Fatal("event type error")
    67  				continue
    68  			}
    69  			gotevents = append(gotevents, ev)
    70  		case <-timeout:
    71  			if !reflect.DeepEqual(gotevents, wantEvents) {
    72  				t.Fatalf("mismatch for test lan discover got %s want %s", spew.Sdump(gotevents), spew.Sdump(wantEvents))
    73  			}
    74  			return
    75  		}
    76  	}
    77  }