test: try WANT_HAVE and WANT_BLOCK in TestPeerBlockFilter

internal/decision/engine_test.go

@@ -1156,42 +1156,113 @@ func TestPeerBlockFilter(t *testing.T) {
 	)
 	e.StartWorkers(ctx, process.WithTeardown(func() error { return nil }))
 
-	// Create wants requests
-	for _, peerID := range peerIDs {
-		partnerWantBlocks(e, keys, peerID)
+	// Setup the test
+	type testCaseEntry struct {
+		peerIndex    int
+		wantBlks     string
+		wantHaves    string
+		sendDontHave bool
 	}
 
-	// check that outgoing messages are sent with the correct content
-	checkPeer := func(peerIndex int, expectedBlocks []blocks.Block) {
-		next := <-e.Outbox()
-		envelope := <-next
+	type testCaseExp struct {
+		blks      string
+		haves     string
+		dontHaves string
+	}
 
-		peerID := peerIDs[peerIndex]
-		responseBlocks := envelope.Message.Blocks()
+	type testCase struct {
+		only bool
+		wl   testCaseEntry
+		exp  testCaseExp
+	}
 
-		if peerID != envelope.Peer {
-			t.Errorf("(Peer%v) expected message for peer ID %#v but instead got message for peer ID %#v", peerIndex, peerID, envelope.Peer)
-		}
+	testCases := []testCase{
+		{
+			wl: testCaseEntry{
+				peerIndex:    0,
+				wantBlks:     "a",
+				sendDontHave: true,
+			},
+			exp: testCaseExp{
+				blks: "a",
+			},
+		},
+		{
+			wl: testCaseEntry{
+				peerIndex:    0,
+				wantHaves:    "b01",
+				sendDontHave: true,
+			},
+			exp: testCaseExp{
+				haves:     "b",
+				dontHaves: "01",
+			},
+		},
+		{
+			wl: testCaseEntry{
+				peerIndex:    1,
+				wantHaves:    "ab",
+				wantBlks:     "c01",
+				sendDontHave: true,
+			},
+			exp: testCaseExp{
+				haves:     "b",
+				blks:      "c",
+				dontHaves: "a01",
+			},
+		},
+		{
+			wl: testCaseEntry{
+				peerIndex:    2,
+				wantHaves:    "ab",
+				wantBlks:     "c01",
+				sendDontHave: true,
+			},
+			exp: testCaseExp{
+				haves:     "",
+				blks:      "c",
+				dontHaves: "ab01",
+			},
+		},
+	}
 
-		if len(responseBlocks) != len(expectedBlocks) {
-			t.Errorf("(Peer%v) expected %v block in response but instead got %v", peerIndex, len(expectedBlocks), len(responseBlocks))
+	var onlyTestCases []testCase
+	for _, testCase := range testCases {
+		if testCase.only {
+			onlyTestCases = append(onlyTestCases, testCase)
 		}
+	}
+	if len(onlyTestCases) > 0 {
+		testCases = onlyTestCases
+	}
 
-		responseBlockSet := make(map[cid.Cid]bool)
-		for _, b := range responseBlocks {
-			responseBlockSet[b.Cid()] = true
-		}
+	for i, testCase := range testCases {
+		// Create wants requests
+		wl := testCase.wl
 
-		for _, b := range expectedBlocks {
-			if !responseBlockSet[b.Cid()] {
-				t.Errorf("(Peer%v) expected block with CID %v", peerIndex, b.Cid())
-			}
+		t.Logf("test case %v: Peer%v / want-blocks '%s' / want-haves '%s' / sendDontHave %t",
+			i, wl.peerIndex, wl.wantBlks, wl.wantHaves, wl.sendDontHave)
+
+		wantBlks := strings.Split(wl.wantBlks, "")
+		wantHaves := strings.Split(wl.wantHaves, "")
+
+		partnerWantBlocksHaves(e, wantBlks, wantHaves, wl.sendDontHave, peerIDs[wl.peerIndex])
+
+		// Check result
+		exp := testCase.exp
+
+		next := <-e.Outbox()
+		envelope := <-next
+
+		expBlks := strings.Split(exp.blks, "")
+		expHaves := strings.Split(exp.haves, "")
+		expDontHaves := strings.Split(exp.dontHaves, "")
+
+		err := checkOutput(t, e, envelope, expBlks, expHaves, expDontHaves)
+		if err != nil {
+			t.Fatal(err)
 		}
 	}
-
-	checkPeer(0, blks[0:3])
-	checkPeer(1, blks[1:3])
-	checkPeer(2, blks[2:3])
 }
 
 func TestTaggingPeers(t *testing.T) {