bitswap: add better tests around wantlist clearing

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>

exchange/bitswap/bitswap_test.go

@@ -334,7 +334,6 @@ func TestDoubleGet(t *testing.T) {
 	blocks := bg.Blocks(1)
 
 	ctx1, cancel1 := context.WithCancel(context.Background())
-
 	blkch1, err := instances[1].Exchange.GetBlocks(ctx1, []key.Key{blocks[0].Key()})
 	if err != nil {
 		t.Fatal(err)
@@ -362,11 +361,15 @@ func TestDoubleGet(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	blk, ok := <-blkch2
-	if !ok {
-		t.Fatal("expected to get the block here")
+	select {
+	case blk, ok := <-blkch2:
+		if !ok {
+			t.Fatal("expected to get the block here")
+		}
+		t.Log(blk)
+	case <-time.After(time.Second * 5):
+		t.Fatal("timed out waiting on block")
 	}
-	t.Log(blk)
 
 	for _, inst := range instances {
 		err := inst.Exchange.Close()
@@ -375,3 +378,68 @@ func TestDoubleGet(t *testing.T) {
 		}
 	}
 }
+
+func TestWantlistCleanup(t *testing.T) {
+	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
+	sg := NewTestSessionGenerator(net)
+	defer sg.Close()
+	bg := blocksutil.NewBlockGenerator()
+
+	instances := sg.Instances(1)[0]
+	bswap := instances.Exchange
+	blocks := bg.Blocks(20)
+
+	var keys []key.Key
+	for _, b := range blocks {
+		keys = append(keys, b.Key())
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*50)
+	defer cancel()
+	_, err := bswap.GetBlock(ctx, keys[0])
+	if err != context.DeadlineExceeded {
+		t.Fatal("shouldnt have fetched any blocks")
+	}
+
+	time.Sleep(time.Millisecond * 50)
+
+	if len(bswap.GetWantlist()) > 0 {
+		t.Fatal("should not have anyting in wantlist")
+	}
+
+	ctx, cancel = context.WithTimeout(context.Background(), time.Millisecond*50)
+	defer cancel()
+	_, err = bswap.GetBlocks(ctx, keys[:10])
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	<-ctx.Done()
+	time.Sleep(time.Millisecond * 50)
+
+	if len(bswap.GetWantlist()) > 0 {
+		t.Fatal("should not have anyting in wantlist")
+	}
+
+	_, err = bswap.GetBlocks(context.Background(), keys[:1])
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel = context.WithCancel(context.Background())
+	_, err = bswap.GetBlocks(ctx, keys[10:])
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	time.Sleep(time.Millisecond * 50)
+	if len(bswap.GetWantlist()) != 11 {
+		t.Fatal("should have 11 keys in wantlist")
+	}
+
+	cancel()
+	time.Sleep(time.Millisecond * 50)
+	if !(len(bswap.GetWantlist()) == 1 && bswap.GetWantlist()[0] == keys[0]) {
+		t.Fatal("should only have keys[0] in wantlist")
+	}
+}