testutil.go 3.31 KB
Newer Older
1 2 3
package testutil

import (
Jakub Sztandera's avatar
Jakub Sztandera committed
4
	"math/rand"
5

6
	bsmsg "github.com/ipfs/go-bitswap/message"
7
	bssd "github.com/ipfs/go-bitswap/sessiondata"
8
	"github.com/ipfs/go-bitswap/wantlist"
9
	blocks "github.com/ipfs/go-block-format"
10 11
	cid "github.com/ipfs/go-cid"
	blocksutil "github.com/ipfs/go-ipfs-blocksutil"
Raúl Kripalani's avatar
Raúl Kripalani committed
12
	peer "github.com/libp2p/go-libp2p-core/peer"
13 14 15 16
)

var blockGenerator = blocksutil.NewBlockGenerator()
var prioritySeq int
17 18 19 20 21

// GenerateBlocksOfSize generates a series of blocks of the given byte size
func GenerateBlocksOfSize(n int, size int64) []blocks.Block {
	generatedBlocks := make([]blocks.Block, 0, n)
	for i := 0; i < n; i++ {
Jakub Sztandera's avatar
Jakub Sztandera committed
22
		// rand.Read never errors
Steven Allen's avatar
Steven Allen committed
23
		buf := make([]byte, size)
Jakub Sztandera's avatar
Jakub Sztandera committed
24 25
		rand.Read(buf)
		b := blocks.NewBlock(buf)
26 27 28 29 30
		generatedBlocks = append(generatedBlocks, b)

	}
	return generatedBlocks
}
31

32
// GenerateCids produces n content identifiers.
33 34 35 36 37 38 39 40 41
func GenerateCids(n int) []cid.Cid {
	cids := make([]cid.Cid, 0, n)
	for i := 0; i < n; i++ {
		c := blockGenerator.Next().Cid()
		cids = append(cids, c)
	}
	return cids
}

42
// GenerateWantlist makes a populated wantlist.
43 44
func GenerateWantlist(n int, ses uint64) *wantlist.SessionTrackedWantlist {
	wl := wantlist.NewSessionTrackedWantlist()
45 46 47 48 49 50 51 52
	for i := 0; i < n; i++ {
		prioritySeq++
		entry := wantlist.NewRefEntry(blockGenerator.Next().Cid(), prioritySeq)
		wl.AddEntry(entry, ses)
	}
	return wl
}

53
// GenerateMessageEntries makes fake bitswap message entries.
54 55
func GenerateMessageEntries(n int, isCancel bool) []bsmsg.Entry {
	bsmsgs := make([]bsmsg.Entry, 0, n)
56 57
	for i := 0; i < n; i++ {
		prioritySeq++
58
		msg := bsmsg.Entry{
59 60 61 62 63 64 65 66 67 68
			Entry:  wantlist.NewRefEntry(blockGenerator.Next().Cid(), prioritySeq),
			Cancel: isCancel,
		}
		bsmsgs = append(bsmsgs, msg)
	}
	return bsmsgs
}

var peerSeq int

69
// GeneratePeers creates n peer ids.
70 71 72 73 74 75 76 77 78 79
func GeneratePeers(n int) []peer.ID {
	peerIds := make([]peer.ID, 0, n)
	for i := 0; i < n; i++ {
		peerSeq++
		p := peer.ID(peerSeq)
		peerIds = append(peerIds, p)
	}
	return peerIds
}

80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97
// GenerateOptimizedPeers creates n peer ids,
// with optimization fall off up to optCount, curveFunc to scale it
func GenerateOptimizedPeers(n int, optCount int, curveFunc func(float64) float64) []bssd.OptimizedPeer {
	peers := GeneratePeers(n)
	optimizedPeers := make([]bssd.OptimizedPeer, 0, n)
	for i, peer := range peers {
		var optimizationRating float64
		if i <= optCount {
			optimizationRating = 1.0 - float64(i)/float64(optCount)
		} else {
			optimizationRating = 0.0
		}
		optimizationRating = curveFunc(optimizationRating)
		optimizedPeers = append(optimizedPeers, bssd.OptimizedPeer{Peer: peer, OptimizationRating: optimizationRating})
	}
	return optimizedPeers
}

98 99
var nextSession uint64

100
// GenerateSessionID make a unit session identifier.
101 102 103 104 105
func GenerateSessionID() uint64 {
	nextSession++
	return uint64(nextSession)
}

106
// ContainsPeer returns true if a peer is found n a list of peers.
107 108 109 110 111 112 113 114
func ContainsPeer(peers []peer.ID, p peer.ID) bool {
	for _, n := range peers {
		if p == n {
			return true
		}
	}
	return false
}
hannahhoward's avatar
hannahhoward committed
115

116 117 118 119 120
// IndexOf returns the index of a given cid in an array of blocks
func IndexOf(blks []blocks.Block, c cid.Cid) int {
	for i, n := range blks {
		if n.Cid() == c {
			return i
hannahhoward's avatar
hannahhoward committed
121 122
		}
	}
123 124 125 126 127 128
	return -1
}

// ContainsBlock returns true if a block is found n a list of blocks
func ContainsBlock(blks []blocks.Block, block blocks.Block) bool {
	return IndexOf(blks, block.Cid()) != -1
hannahhoward's avatar
hannahhoward committed
129
}