github.com/annchain/OG@v0.0.9/og/peer_test.go (about)

     1  // Copyright © 2019 Annchain Authors <EMAIL ADDRESS>
     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 og
    15  
    16  import (
    17  	"fmt"
    18  	"github.com/annchain/OG/og/types/archive"
    19  
    20  	"github.com/annchain/OG/p2p"
    21  	"github.com/deckarep/golang-set"
    22  	"testing"
    23  )
    24  
    25  func TestPeerSet_GetRandomPeers(t *testing.T) {
    26  
    27  	set := newPeerSet()
    28  	for i := 0; i < 10; i++ {
    29  		rawPeer := p2p.Peer{}
    30  		p := &peer{
    31  			Peer:      &rawPeer,
    32  			version:   1,
    33  			id:        fmt.Sprintf("%d", i),
    34  			knownMsg:  mapset.NewSet(),
    35  			queuedMsg: make(chan []*types, maxqueuedMsg),
    36  			term:      make(chan struct{}),
    37  		}
    38  		set.Register(p)
    39  	}
    40  	fmt.Println("len peers", set.Len())
    41  	peers := set.GetRandomPeers(3)
    42  	if len(peers) != 3 {
    43  		t.Fatalf("peers size mismatch, wanted 3 ,got %d ,peers %v ", len(peers), peers)
    44  	}
    45  	peers = set.GetPeers(nil, 3)
    46  	if len(peers) != 0 {
    47  		t.Fatalf("peers size mismatch, wanted 3 ,got %d ,peers %v", len(peers), peers)
    48  	}
    49  
    50  }
    51  
    52  func TestPeer_MarkMessage(t *testing.T) {
    53  	rawPeer := p2p.Peer{}
    54  	p := &peer{
    55  		Peer:      &rawPeer,
    56  		version:   1,
    57  		id:        "111",
    58  		knownMsg:  mapset.NewSet(),
    59  		queuedMsg: make(chan []*types, maxqueuedMsg),
    60  		term:      make(chan struct{}),
    61  	}
    62  	var msgs []types
    63  	for i := 0; i < 100; i++ {
    64  		msg := types{Message: &p2p_message.MessageNewTx{RawTx: archive.RandomTx().RawTx()}, MessageType: message_archive.MessageTypeNewTx}
    65  		msg.Marshal()
    66  		msg.CalculateHash()
    67  		msgs = append(msgs, msg)
    68  		p.MarkMessage(msg.MessageType, *msg.Hash)
    69  	}
    70  
    71  	for i, val := range msgs {
    72  		key := val.MsgKey()
    73  		if !p.knownMsg.Contains(key) {
    74  			t.Fatal(i, val, key)
    75  		}
    76  	}
    77  
    78  }