github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/xorlayer/xtm_test.go (about)

     1  package xorlayer
     2  
     3  import (
     4  	"reflect"
     5  	"testing"
     6  
     7  	"sort"
     8  
     9  	"math/rand"
    10  
    11  	"gotest.tools/assert"
    12  )
    13  
    14  func TestXTM(t *testing.T) {
    15  	nodeID := NodeID(1)
    16  	k, h := 4, 1
    17  	x := NewXTM(k, h, nodeID, nil)
    18  
    19  	assert.Assert(t, x.getBucketIdx(nodeID) == bitSize)
    20  
    21  	total := 10000
    22  	cookie := uint64(0)
    23  	neighbours := make(map[NodeID]bool)
    24  	for i := 0; i < total; i++ {
    25  		n := NodeID(rand.Uint64())
    26  		if n == nodeID {
    27  			continue
    28  		}
    29  		if neighbours[n] {
    30  			continue
    31  		}
    32  		neighbours[n] = true
    33  		x.AddNeighbour(n, cookie)
    34  	}
    35  
    36  	kclosest := x.KClosest(nodeID)
    37  	sort.Slice(kclosest, func(i, j int) bool {
    38  		return kclosest[i] < kclosest[j]
    39  	})
    40  
    41  	neighbourSlice := make([]NodeID, 0, len(neighbours))
    42  	for n := range neighbours {
    43  		neighbourSlice = append(neighbourSlice, n)
    44  	}
    45  	sort.Slice(neighbourSlice, func(i, j int) bool {
    46  		return neighbourSlice[i]^nodeID < neighbourSlice[j]^nodeID
    47  	})
    48  
    49  	expected := []NodeID{nodeID}
    50  	expected = append(expected, neighbourSlice[0:k-1]...)
    51  	assert.Assert(t, reflect.DeepEqual(kclosest, expected))
    52  }