rebased on announce_tests

cmd/devp2p/internal/ethtest/suite.go

@@ -38,6 +38,8 @@ var pretty = spew.ConfigState{
 	SortKeys:                true,
 }
 
+var timeout = 20 * time.Second
+
 // Suite represents a structure used to test the eth
 // protocol of a node(s).
 type Suite struct {
@@ -120,7 +122,6 @@ func (s *Suite) TestMaliciousStatus(t *utesting.T) {
 	default:
 		t.Fatalf("expected status, got: %#v ", msg)
 	}
-	timeout := 20 * time.Second
 	// wait for disconnect
 	switch msg := conn.ReadAndServe(s.chain, timeout).(type) {
 	case *Disconnect:
@@ -156,7 +157,6 @@ func (s *Suite) TestGetBlockHeaders(t *utesting.T) {
 		t.Fatalf("could not write to connection: %v", err)
 	}
 
-	timeout := 20 * time.Second
 	switch msg := conn.ReadAndServe(s.chain, timeout).(type) {
 	case *BlockHeaders:
 		headers := msg
@@ -186,7 +186,6 @@ func (s *Suite) TestGetBlockBodies(t *utesting.T) {
 		t.Fatalf("could not write to connection: %v", err)
 	}
 
-	timeout := 20 * time.Second
 	switch msg := conn.ReadAndServe(s.chain, timeout).(type) {
 	case *BlockBodies:
 		t.Logf("received %d block bodies", len(*msg))
@@ -318,7 +317,6 @@ func (s *Suite) TestLargeAnnounce(t *utesting.T) {
 			t.Fatalf("could not write to connection: %v", err)
 		}
 		// Invalid announcement, check that peer disconnected
-		timeout := 20 * time.Second
 		switch msg := sendConn.ReadAndServe(s.chain, timeout).(type) {
 		case *Disconnect:
 		case *Error:

cmd/devp2p/internal/ethtest/transaction.go

@@ -36,16 +36,17 @@ func sendSuccessfulTx(t *utesting.T, s *Suite, tx *types.Transaction) {
 		t.Fatal(err)
 	}
 	// Wait for the transaction announcement
-	rawTxMsg, err := recvConn.waitForMessage(Transactions{})
-	if err != nil {
-		t.Fatalf("waiting for transaction propagation failed: %v", err)
-	}
-	recTxs := *rawTxMsg.(*Transactions)
-	if len(recTxs) != 1 {
-		t.Fatalf("received transactions do not match send: %v", recTxs)
-	}
-	if tx.Hash() != recTxs[0].Hash() {
-		t.Fatalf("received transactions do not match send: got %v want %v", tx, recTxs)
+	switch msg := recvConn.ReadAndServe(s.chain, timeout).(type) {
+	case *Transactions:
+		recTxs := *msg
+		if len(recTxs) != 1 {
+			t.Fatalf("received transactions do not match send: %v", recTxs)
+		}
+		if tx.Hash() != recTxs[0].Hash() {
+			t.Fatalf("received transactions do not match send: got %v want %v", tx, recTxs)
+		}
+	default:
+		t.Fatalf("unexpected: %s", pretty.Sdump(msg))
 	}
 }
 
@@ -56,9 +57,11 @@ func sendFailingTx(t *utesting.T, s *Suite, tx *types.Transaction) {
 		t.Fatal(err)
 	}
 	// Wait for the transaction announcement
-	_, err := recvConn.waitForMessage(Transactions{})
-	if err == nil {
-		t.Fatalf("received a transaction, but wanted none", err)
+	switch msg := recvConn.ReadAndServe(s.chain, timeout).(type) {
+	case *Transactions:
+		t.Fatalf("Received unexpected transaction announcement: %v", msg)
+	default:
+		t.Fatalf("unexpected: %s", pretty.Sdump(msg))
 	}
 }