message_test.go 4.07 KB
Newer Older
1 2 3 4 5 6
package message

import (
	"bytes"
	"testing"

7
	proto "gx/ipfs/QmZ4Qi3GaRbjcx28Sme5eMH7RQjGkt8wHxt2a65oLaeFEV/gogo-protobuf/proto"
Jeromy's avatar
Jeromy committed
8

9
	blocks "github.com/ipfs/go-ipfs/blocks"
10
	pb "github.com/ipfs/go-ipfs/exchange/bitswap/message/pb"
11 12
	cid "gx/ipfs/QmakyCk6Vnn16WEKjbkxieZmM2YLTzkFWizbmGowoYPjro/go-cid"
	u "gx/ipfs/Qmb912gdngC1UWwTkhuW8knyRbcWeu5kqkxBpveLmW8bSr/go-ipfs-util"
13 14
)

15 16 17 18
func mkFakeCid(s string) *cid.Cid {
	return cid.NewCidV0(u.Hash([]byte(s)))
}

19
func TestAppendWanted(t *testing.T) {
20
	str := mkFakeCid("foo")
21
	m := New(true)
22
	m.AddEntry(str, 1)
23

Jeromy's avatar
Jeromy committed
24
	if !wantlistContains(m.ToProto().GetWantlist(), str) {
25 26
		t.Fail()
	}
Jeromy's avatar
Jeromy committed
27
	m.ToProto().GetWantlist().GetEntries()
28 29 30
}

func TestNewMessageFromProto(t *testing.T) {
31
	str := mkFakeCid("a_key")
32
	protoMessage := new(pb.Message)
Jeromy's avatar
Jeromy committed
33 34
	protoMessage.Wantlist = new(pb.Message_Wantlist)
	protoMessage.Wantlist.Entries = []*pb.Message_Wantlist_Entry{
35
		{Block: proto.String(str.KeyString())},
Jeromy's avatar
Jeromy committed
36 37
	}
	if !wantlistContains(protoMessage.Wantlist, str) {
38 39
		t.Fail()
	}
40 41 42 43 44
	m, err := newMessageFromProto(*protoMessage)
	if err != nil {
		t.Fatal(err)
	}

Jeromy's avatar
Jeromy committed
45
	if !wantlistContains(m.ToProto().GetWantlist(), str) {
46 47 48 49 50 51 52 53 54 55
		t.Fail()
	}
}

func TestAppendBlock(t *testing.T) {

	strs := make([]string, 2)
	strs = append(strs, "Celeritas")
	strs = append(strs, "Incendia")

56
	m := New(true)
57
	for _, str := range strs {
58
		block := blocks.NewBlock([]byte(str))
Jeromy's avatar
Jeromy committed
59
		m.AddBlock(block)
60 61 62 63 64 65 66 67 68 69 70
	}

	// assert strings are in proto message
	for _, blockbytes := range m.ToProto().GetBlocks() {
		s := bytes.NewBuffer(blockbytes).String()
		if !contains(strs, s) {
			t.Fail()
		}
	}
}

71
func TestWantlist(t *testing.T) {
72
	keystrs := []*cid.Cid{mkFakeCid("foo"), mkFakeCid("bar"), mkFakeCid("baz"), mkFakeCid("bat")}
73
	m := New(true)
74
	for _, s := range keystrs {
75
		m.AddEntry(s, 1)
76 77 78 79 80 81 82
	}
	exported := m.Wantlist()

	for _, k := range exported {
		present := false
		for _, s := range keystrs {

83
			if s.Equals(k.Cid) {
84 85 86 87
				present = true
			}
		}
		if !present {
88
			t.Logf("%v isn't in original list", k.Cid)
89 90 91 92 93
			t.Fail()
		}
	}
}

94
func TestCopyProtoByValue(t *testing.T) {
95
	str := mkFakeCid("foo")
96
	m := New(true)
97
	protoBeforeAppend := m.ToProto()
98
	m.AddEntry(str, 1)
Jeromy's avatar
Jeromy committed
99
	if wantlistContains(protoBeforeAppend.GetWantlist(), str) {
100 101 102 103
		t.Fail()
	}
}

104
func TestToNetFromNetPreservesWantList(t *testing.T) {
105
	original := New(true)
106 107 108 109 110
	original.AddEntry(mkFakeCid("M"), 1)
	original.AddEntry(mkFakeCid("B"), 1)
	original.AddEntry(mkFakeCid("D"), 1)
	original.AddEntry(mkFakeCid("T"), 1)
	original.AddEntry(mkFakeCid("F"), 1)
111

112 113
	buf := new(bytes.Buffer)
	if err := original.ToNet(buf); err != nil {
114 115 116
		t.Fatal(err)
	}

117
	copied, err := FromNet(buf)
118 119 120 121
	if err != nil {
		t.Fatal(err)
	}

122
	keys := make(map[string]bool)
123
	for _, k := range copied.Wantlist() {
124
		keys[k.Cid.KeyString()] = true
125 126 127
	}

	for _, k := range original.Wantlist() {
128
		if _, ok := keys[k.Cid.KeyString()]; !ok {
129 130 131 132 133 134 135
			t.Fatalf("Key Missing: \"%v\"", k)
		}
	}
}

func TestToAndFromNetMessage(t *testing.T) {

136
	original := New(true)
Jeromy's avatar
Jeromy committed
137 138 139 140
	original.AddBlock(blocks.NewBlock([]byte("W")))
	original.AddBlock(blocks.NewBlock([]byte("E")))
	original.AddBlock(blocks.NewBlock([]byte("F")))
	original.AddBlock(blocks.NewBlock([]byte("M")))
141

142 143
	buf := new(bytes.Buffer)
	if err := original.ToNet(buf); err != nil {
144 145 146
		t.Fatal(err)
	}

147
	m2, err := FromNet(buf)
148 149 150 151
	if err != nil {
		t.Fatal(err)
	}

152
	keys := make(map[string]bool)
153
	for _, b := range m2.Blocks() {
154
		keys[b.Cid().KeyString()] = true
155 156 157
	}

	for _, b := range original.Blocks() {
158
		if _, ok := keys[b.Cid().KeyString()]; !ok {
159 160 161 162 163
			t.Fail()
		}
	}
}

164
func wantlistContains(wantlist *pb.Message_Wantlist, c *cid.Cid) bool {
Jeromy's avatar
Jeromy committed
165
	for _, e := range wantlist.GetEntries() {
166
		if e.GetBlock() == c.KeyString() {
Jeromy's avatar
Jeromy committed
167 168 169 170 171 172 173 174 175
			return true
		}
	}
	return false
}

func contains(strs []string, x string) bool {
	for _, s := range strs {
		if s == x {
176 177 178 179 180
			return true
		}
	}
	return false
}
181 182 183

func TestDuplicates(t *testing.T) {
	b := blocks.NewBlock([]byte("foo"))
184
	msg := New(true)
185

186 187
	msg.AddEntry(b.Cid(), 1)
	msg.AddEntry(b.Cid(), 1)
188 189 190 191
	if len(msg.Wantlist()) != 1 {
		t.Fatal("Duplicate in BitSwapMessage")
	}

Jeromy's avatar
Jeromy committed
192 193
	msg.AddBlock(b)
	msg.AddBlock(b)
194 195 196 197
	if len(msg.Blocks()) != 1 {
		t.Fatal("Duplicate in BitSwapMessage")
	}
}