[NC-2056] Remove vertx from discovery tests (#539)

ethereum/eth/src/test/java/tech/pegasys/pantheon/ethereum/eth/transactions/TestNode.java

@@ -122,7 +122,7 @@ public class TestNode implements Closeable {
             .metricsSystem(new NoOpMetricsSystem())
             .build();
     network = networkRunner.getNetwork();
-    this.port = network.getSelf().getPort();
+    this.port = network.getLocalPeerInfo().getPort();
     network.subscribeDisconnect(
         (connection, reason, initiatedByPeer) -> disconnections.put(connection, reason));
 

ethereum/mock-p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/testing/MockNetwork.java

@@ -180,7 +180,7 @@ public final class MockNetwork {
     public void close() {}
 
     @Override
-    public PeerInfo getSelf() {
+    public PeerInfo getLocalPeerInfo() {
       return new PeerInfo(
           5, self.getId().toString(), new ArrayList<>(capabilities), 0, self.getId());
     }

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/api/P2PNetwork.java

@@ -78,7 +78,7 @@ public interface P2PNetwork extends Closeable, Runnable {
    *
    * @return the PeerInfo for this node.
    */
-  PeerInfo getSelf();
+  PeerInfo getLocalPeerInfo();
 
   /**
    * Checks if the node is listening for network connections

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryAgent.java

@@ -13,9 +13,9 @@
 package tech.pegasys.pantheon.ethereum.p2p.discovery;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
-import static tech.pegasys.pantheon.util.Preconditions.checkGuard;
 import static tech.pegasys.pantheon.util.bytes.BytesValue.wrapBuffer;
 
 import tech.pegasys.pantheon.crypto.SECP256K1;
@@ -23,28 +23,26 @@ import tech.pegasys.pantheon.ethereum.p2p.api.DisconnectCallback;
 import tech.pegasys.pantheon.ethereum.p2p.api.PeerConnection;
 import tech.pegasys.pantheon.ethereum.p2p.config.DiscoveryConfiguration;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryEvent.PeerBondedEvent;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryEvent.PeerDroppedEvent;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketData;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketType;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerDiscoveryController;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerRequirement;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerTable;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PingPacketData;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.TimerUtil;
 import tech.pegasys.pantheon.ethereum.p2p.peers.DefaultPeerId;
+import tech.pegasys.pantheon.ethereum.p2p.peers.Endpoint;
 import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
 import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
 import tech.pegasys.pantheon.ethereum.p2p.wire.messages.DisconnectMessage;
 import tech.pegasys.pantheon.util.NetworkUtility;
+import tech.pegasys.pantheon.util.Subscribers;
 import tech.pegasys.pantheon.util.bytes.BytesValue;
 
-import java.io.IOException;
-import java.net.BindException;
 import java.net.InetSocketAddress;
-import java.net.SocketException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
@@ -53,282 +51,201 @@ import java.util.stream.Collectors;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.net.InetAddresses;
-import io.vertx.core.Vertx;
-import io.vertx.core.datagram.DatagramPacket;
-import io.vertx.core.datagram.DatagramSocket;
-import io.vertx.core.datagram.DatagramSocketOptions;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 /**
- * The peer discovery agent is the network component that sends and receives messages peer discovery
+ * The peer discovery agent is the network component that sends and receives peer discovery messages
- * messages via UDP. It exposes methods for the {@link PeerDiscoveryController} to dispatch outbound
+ * via UDP.
- * messages too.
- *
- * <h3>How do the peer table and the discovery agent interact with one another?</h3>
- *
- * <ul>
- *   <li>The agent acts like the transport layer, receiving messages from the wire and exposing
- *       methods for the peer table to send packets too.
- *   <li>The table stores and indexes peers in a Kademlia k-bucket table with 256 bins (where bin 0
- *       is not used as it's us, i.e. distance 0 == us). It reacts to messages based on its internal
- *       state. It uses the agent whenever it needs to dispatch a message.
- * </ul>
- *
- * <h3>The flow</h3>
- *
- * <ol>
- *   <li>The discovery agent dispatches all incoming messages that were properly decoded and whose
- *       hash integrity check passes to the peer table.
- *   <li>The peer table decides whether to store the Peer, change its state, send other messages,
- *       etc. based on its internal state.
- *   <li>The agent attaches a callback to the call to the Peer Table. When the Peer Table has
- *       processed the message, it'll perform a callback passing in an Optional which is populated
- *       if we recognised the Peer, and empty if we did not.
- *   <li>The agent reacts to specific messages (PING-&gt;PONG, FIND_NEIGHBORS-&gt;NEIGHBORS), if the
- *       Peer was recognised. Why doesn't the table send these messages itself? Because they don't
- *       affect the state machine of the Peer, and the table is only concerned with storing peers,
- *       keeping them alive and tracking their state. It is not bothered to service requests.
- * </ol>
  */
-public class PeerDiscoveryAgent implements DisconnectCallback {
+public abstract class PeerDiscoveryAgent implements DisconnectCallback {
-  private static final Logger LOG = LogManager.getLogger();
+  protected static final Logger LOG = LogManager.getLogger();
 
   // The devp2p specification says only accept packets up to 1280, but some
   // clients ignore that, so we add in a little extra padding.
   private static final int MAX_PACKET_SIZE_BYTES = 1600;
   private static final long PEER_REFRESH_INTERVAL_MS = MILLISECONDS.convert(30, TimeUnit.MINUTES);
-  private final Vertx vertx;
+
+  protected final List<DiscoveryPeer> bootstrapPeers;
+  private final PeerRequirement peerRequirement;
+  private final PeerBlacklist peerBlacklist;
+  private final NodeWhitelistController nodeWhitelistController;
   /* The peer controller, which takes care of the state machine of peers. */
-  private final PeerDiscoveryController controller;
+  protected Optional<PeerDiscoveryController> controller = Optional.empty();
+
   /* The keypair used to sign messages. */
-  private final SECP256K1.KeyPair keyPair;
+  protected final SECP256K1.KeyPair keyPair;
+  private final BytesValue id;
   private final PeerTable peerTable;
-  private final DiscoveryConfiguration config;
+  protected final DiscoveryConfiguration config;
 
   /* This is the {@link tech.pegasys.pantheon.ethereum.p2p.Peer} object holding who we are. */
   private DiscoveryPeer advertisedPeer;
-  /* The vert.x UDP socket. */
+  private InetSocketAddress localAddress;
-  private DatagramSocket socket;
 
   /* Is discovery enabled? */
   private boolean isActive = false;
+  private final Subscribers<Consumer<PeerBondedEvent>> peerBondedObservers = new Subscribers<>();
 
   public PeerDiscoveryAgent(
-      final Vertx vertx,
       final SECP256K1.KeyPair keyPair,
       final DiscoveryConfiguration config,
       final PeerRequirement peerRequirement,
       final PeerBlacklist peerBlacklist,
       final NodeWhitelistController nodeWhitelistController) {
-    checkArgument(vertx != null, "vertx instance cannot be null");
     checkArgument(keyPair != null, "keypair cannot be null");
     checkArgument(config != null, "provided configuration cannot be null");
 
     validateConfiguration(config);
 
-    final List<DiscoveryPeer> bootstrapPeers =
+    this.peerRequirement = peerRequirement;
+    this.peerBlacklist = peerBlacklist;
+    this.nodeWhitelistController = nodeWhitelistController;
+    this.bootstrapPeers =
         config.getBootstrapPeers().stream().map(DiscoveryPeer::new).collect(Collectors.toList());
 
-    this.vertx = vertx;
     this.config = config;
     this.keyPair = keyPair;
     this.peerTable = new PeerTable(keyPair.getPublicKey().getEncodedBytes(), 16);
-    this.controller =
+
-        new PeerDiscoveryController(
+    id = keyPair.getPublicKey().getEncodedBytes();
-            vertx,
-            this,
-            peerTable,
-            bootstrapPeers,
-            PEER_REFRESH_INTERVAL_MS,
-            peerRequirement,
-            peerBlacklist,
-            nodeWhitelistController);
   }
 
-  public CompletableFuture<?> start(final int tcpPort) {
+  protected abstract TimerUtil createTimer();
-    final CompletableFuture<?> completion = new CompletableFuture<>();
+
+  protected abstract CompletableFuture<InetSocketAddress> listenForConnections();
+
+  protected abstract CompletableFuture<Void> sendOutgoingPacket(
+      final DiscoveryPeer peer, final Packet packet);
+
+  public abstract CompletableFuture<?> stop();
+
+  public CompletableFuture<?> start() {
+    final CompletableFuture<?> future = new CompletableFuture<>();
     if (config.isActive()) {
       final String host = config.getBindHost();
       final int port = config.getBindPort();
       LOG.info("Starting peer discovery agent on host={}, port={}", host, port);
 
-      vertx
+      listenForConnections()
-          .createDatagramSocket(
+          .thenAccept(
-              new DatagramSocketOptions().setIpV6(NetworkUtility.isIPv6Available()))
+              (InetSocketAddress localAddress) -> {
-          .listen(
+                // Once listener is set up, finish initializing
-              port,
+                this.localAddress = localAddress;
-              host,
+                advertisedPeer =
-              res -> {
+                    new DiscoveryPeer(
-                if (res.failed()) {
+                        id,
-                  Throwable cause = res.cause();
+                        config.getAdvertisedHost(),
-                  LOG.error("An exception occurred when starting the peer discovery agent", cause);
+                        localAddress.getPort(),
-
+                        localAddress.getPort());
-                  if (cause instanceof BindException || cause instanceof SocketException) {
+                isActive = true;
-                    cause =
+                startController();
-                        new PeerDiscoveryServiceException(
+              })
-                            String.format(
+          .whenComplete(
-                                "Failed to bind Ethereum UDP discovery listener to %s:%d: %s",
+              (res, err) -> {
-                                host, port, cause.getMessage()));
+                // Finalize future
-                  }
+                if (err != null) {
-                  completion.completeExceptionally(cause);
+                  future.completeExceptionally(err);
-                  return;
+                } else {
+                  future.complete(null);
                 }
-                initialize(res.result(), res.result().localAddress().port());
-                this.isActive = true;
-                completion.complete(null);
               });
     } else {
       this.isActive = false;
-      completion.complete(null);
+      future.complete(null);
-    }
-    return completion;
-  }
-
-  public CompletableFuture<?> stop() {
-    if (socket == null) {
-      return CompletableFuture.completedFuture(null);
-    }
-
-    final CompletableFuture<?> completion = new CompletableFuture<>();
-    socket.close(
-        ar -> {
-          if (ar.succeeded()) {
-            controller.stop();
-            socket = null;
-            completion.complete(null);
-          } else {
-            completion.completeExceptionally(ar.cause());
     }
-        });
+    return future;
-    return completion;
   }
 
-  private void initialize(final DatagramSocket socket, final int tcpPort) {
+  private void startController() {
-    this.socket = socket;
+    PeerDiscoveryController controller = createController();
-
+    this.controller = Optional.of(controller);
-    // TODO: when using wildcard hosts (0.0.0.0), we need to handle multiple addresses by selecting
-    // the
-    // correct 'announce' address.
-    final BytesValue id = keyPair.getPublicKey().getEncodedBytes();
-    final String effectiveHost = socket.localAddress().host();
-    final int effectivePort = socket.localAddress().port();
-    advertisedPeer = new DiscoveryPeer(id, config.getAdvertisedHost(), effectivePort, tcpPort);
-
-    LOG.info(
-        "Started peer discovery agent successfully, on effective host={} and port={}",
-        effectiveHost,
-        effectivePort);
-
-    socket.exceptionHandler(this::handleException);
-    socket.handler(this::handlePacket);
     controller.start();
   }
 
-  /**
+  private PeerDiscoveryController createController() {
-   * For uncontrolled exceptions occurring in the packet handlers.
+    return new PeerDiscoveryController(
-   *
+        keyPair,
-   * @param exception the exception that was raised
+        advertisedPeer,
-   */
+        peerTable,
-  private void handleException(final Throwable exception) {
+        bootstrapPeers,
-    if (exception instanceof IOException) {
+        this::handleOutgoingPacket,
-      LOG.debug("Packet handler exception", exception);
+        createTimer(),
-    } else {
+        PEER_REFRESH_INTERVAL_MS,
-      LOG.error("Packet handler exception", exception);
+        peerRequirement,
+        peerBlacklist,
+        nodeWhitelistController,
+        peerBondedObservers);
   }
+
+  protected boolean validatePacketSize(final int packetSize) {
+    return packetSize <= MAX_PACKET_SIZE_BYTES;
   }
 
-  /**
+  protected void handleIncomingPacket(final Endpoint sourceEndpoint, final Packet packet) {
-   * The UDP packet handler. This is the entrypoint for all received datagrams.
+    OptionalInt tcpPort = OptionalInt.empty();
-   *
-   * @param datagram the received datagram.
-   */
-  private void handlePacket(final DatagramPacket datagram) {
-    try {
-      final int length = datagram.data().length();
-      checkGuard(
-          length <= MAX_PACKET_SIZE_BYTES,
-          PeerDiscoveryPacketDecodingException::new,
-          "Packet too large. Actual size (bytes): %s",
-          length);
-
-      // We allow exceptions to bubble up, as they'll be picked up by the exception handler.
-      final Packet packet = Packet.decode(datagram.data());
-
-      OptionalInt fromPort = OptionalInt.empty();
     if (packet.getPacketData(PingPacketData.class).isPresent()) {
       final PingPacketData ping = packet.getPacketData(PingPacketData.class).orElseGet(null);
       if (ping != null && ping.getFrom() != null && ping.getFrom().getTcpPort().isPresent()) {
-          fromPort = ping.getFrom().getTcpPort();
+        tcpPort = ping.getFrom().getTcpPort();
       }
     }
 
-      // Acquire the senders coordinates to build a Peer representation from them.
-      final String addr = datagram.sender().host();
-      final int port = datagram.sender().port();
-
     // Notify the peer controller.
-      final DiscoveryPeer peer = new DiscoveryPeer(packet.getNodeId(), addr, port, fromPort);
+    String host = sourceEndpoint.getHost();
-      controller.onMessage(packet, peer);
+    int port = sourceEndpoint.getUdpPort();
-    } catch (final PeerDiscoveryPacketDecodingException e) {
+    final DiscoveryPeer peer = new DiscoveryPeer(packet.getNodeId(), host, port, tcpPort);
-      LOG.debug("Discarding invalid peer discovery packet", e);
+    controller.ifPresent(c -> c.onMessage(packet, peer));
-    } catch (final Throwable t) {
-      LOG.error("Encountered error while handling packet", t);
-    }
   }
 
   /**
-   * Allows package-private components to dispatch messages to peers. It updates the lastContacted
+   * Send a packet to the given recipient.
-   * timestamp of the {@link DiscoveryPeer}. This method wraps the data in a Packet, calculates its
-   * hash and signs it with our private key.
    *
    * @param peer the recipient
-   * @param type the type of message
+   * @param packet the packet to send
-   * @param data the data packet to send
-   * @return the sent packet
    */
-  public Packet sendPacket(final DiscoveryPeer peer, final PacketType type, final PacketData data) {
+  protected void handleOutgoingPacket(final DiscoveryPeer peer, final Packet packet) {
-    final Packet packet = Packet.create(type, data, keyPair);
     LOG.trace(
         ">>> Sending {} discovery packet to {} ({}): {}",
-        type,
+        packet.getType(),
         peer.getEndpoint(),
         peer.getId().slice(0, 16),
         packet);
 
-    // Update the lastContacted timestamp on the peer if the dispatch succeeds.
+    sendOutgoingPacket(peer, packet)
-    socket.send(
+        .whenComplete(
-        packet.encode(),
+            (res, err) -> {
-        peer.getEndpoint().getUdpPort(),
+              if (err != null) {
-        peer.getEndpoint().getHost(),
-        ar -> {
-          if (ar.failed()) {
                 LOG.warn(
                     "Sending to peer {} failed, packet: {}",
                     peer,
                     wrapBuffer(packet.encode()),
-                ar.cause());
+                    err);
                 return;
               }
-          if (ar.succeeded()) {
               peer.setLastContacted(System.currentTimeMillis());
-          }
             });
-
-    return packet;
   }
 
+  @VisibleForTesting
   public Collection<DiscoveryPeer> getPeers() {
-    return Collections.unmodifiableCollection(controller.getPeers());
+    return controller
+        .map(PeerDiscoveryController::getPeers)
+        .map(Collections::unmodifiableCollection)
+        .orElse(Collections.emptyList());
   }
 
   public DiscoveryPeer getAdvertisedPeer() {
     return advertisedPeer;
   }
 
+  public BytesValue getId() {
+    return id;
+  }
+
   public InetSocketAddress localAddress() {
-    checkState(socket != null, "uninitialized discovery agent");
+    checkState(localAddress != null, "Uninitialized discovery agent");
-    return new InetSocketAddress(socket.localAddress().host(), socket.localAddress().port());
+    return localAddress;
   }
 
   /**
@@ -341,19 +258,8 @@ public class PeerDiscoveryAgent implements DisconnectCallback {
    * @return A unique ID identifying this observer, to that it can be removed later.
    */
   public long observePeerBondedEvents(final Consumer<PeerBondedEvent> observer) {
-    return controller.observePeerBondedEvents(observer);
+    checkNotNull(observer);
-  }
+    return peerBondedObservers.subscribe(observer);
-
-  /**
-   * Adds an observer that will get called when a new peer is dropped from the peer table.
-   *
-   * <p><i>No guarantees are made about the order in which observers are invoked.</i>
-   *
-   * @param observer The observer to call.
-   * @return A unique ID identifying this observer, to that it can be removed later.
-   */
-  public long observePeerDroppedEvents(final Consumer<PeerDroppedEvent> observer) {
-    return controller.observePeerDroppedEvents(observer);
   }
 
   /**
@@ -363,17 +269,7 @@ public class PeerDiscoveryAgent implements DisconnectCallback {
    * @return Whether the observer was located and removed.
    */
   public boolean removePeerBondedObserver(final long observerId) {
-    return controller.removePeerBondedObserver(observerId);
+    return peerBondedObservers.unsubscribe(observerId);
-  }
-
-  /**
-   * Removes an previously added peer dropped observer.
-   *
-   * @param observerId The unique ID identifying the observer to remove.
-   * @return Whether the observer was located and removed.
-   */
-  public boolean removePeerDroppedObserver(final long observerId) {
-    return controller.removePeerDroppedObserver(observerId);
   }
 
   /**
@@ -383,7 +279,7 @@ public class PeerDiscoveryAgent implements DisconnectCallback {
    */
   @VisibleForTesting
   public int getObserverCount() {
-    return controller.observerCount();
+    return peerBondedObservers.getSubscriberCount();
   }
 
   private static void validateConfiguration(final DiscoveryConfiguration config) {

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/discovery/VertxPeerDiscoveryAgent.java

@@ -0,0 +1,196 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.ethereum.p2p.discovery;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import tech.pegasys.pantheon.crypto.SECP256K1.KeyPair;
+import tech.pegasys.pantheon.ethereum.p2p.config.DiscoveryConfiguration;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerDiscoveryController;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerRequirement;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.TimerUtil;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.VertxTimerUtil;
+import tech.pegasys.pantheon.ethereum.p2p.peers.Endpoint;
+import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
+import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
+import tech.pegasys.pantheon.util.NetworkUtility;
+import tech.pegasys.pantheon.util.Preconditions;
+
+import java.io.IOException;
+import java.net.BindException;
+import java.net.InetSocketAddress;
+import java.net.SocketException;
+import java.util.OptionalInt;
+import java.util.concurrent.CompletableFuture;
+
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Vertx;
+import io.vertx.core.datagram.DatagramPacket;
+import io.vertx.core.datagram.DatagramSocket;
+import io.vertx.core.datagram.DatagramSocketOptions;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+public class VertxPeerDiscoveryAgent extends PeerDiscoveryAgent {
+  private static final Logger LOG = LogManager.getLogger();
+
+  private final Vertx vertx;
+  /* The vert.x UDP socket. */
+  private DatagramSocket socket;
+
+  public VertxPeerDiscoveryAgent(
+      final Vertx vertx,
+      final KeyPair keyPair,
+      final DiscoveryConfiguration config,
+      final PeerRequirement peerRequirement,
+      final PeerBlacklist peerBlacklist,
+      final NodeWhitelistController nodeWhitelistController) {
+    super(keyPair, config, peerRequirement, peerBlacklist, nodeWhitelistController);
+    checkArgument(vertx != null, "vertx instance cannot be null");
+    this.vertx = vertx;
+  }
+
+  @Override
+  protected TimerUtil createTimer() {
+    return new VertxTimerUtil(vertx);
+  }
+
+  @Override
+  protected CompletableFuture<InetSocketAddress> listenForConnections() {
+    CompletableFuture<InetSocketAddress> future = new CompletableFuture<>();
+    vertx
+        .createDatagramSocket(new DatagramSocketOptions().setIpV6(NetworkUtility.isIPv6Available()))
+        .listen(
+            config.getBindPort(), config.getBindHost(), res -> handleListenerSetup(res, future));
+    return future;
+  }
+
+  protected void handleListenerSetup(
+      final AsyncResult<DatagramSocket> listenResult,
+      final CompletableFuture<InetSocketAddress> addressFuture) {
+    if (listenResult.failed()) {
+      Throwable cause = listenResult.cause();
+      LOG.error("An exception occurred when starting the peer discovery agent", cause);
+
+      if (cause instanceof BindException || cause instanceof SocketException) {
+        cause =
+            new PeerDiscoveryServiceException(
+                String.format(
+                    "Failed to bind Ethereum UDP discovery listener to %s:%d: %s",
+                    config.getBindHost(), config.getBindPort(), cause.getMessage()));
+      }
+      addressFuture.completeExceptionally(cause);
+      return;
+    }
+
+    this.socket = listenResult.result();
+
+    // TODO: when using wildcard hosts (0.0.0.0), we need to handle multiple addresses by
+    // selecting
+    // the correct 'announce' address.
+    final String effectiveHost = socket.localAddress().host();
+    final int effectivePort = socket.localAddress().port();
+
+    LOG.info(
+        "Started peer discovery agent successfully, on effective host={} and port={}",
+        effectiveHost,
+        effectivePort);
+
+    socket.exceptionHandler(this::handleException);
+    socket.handler(this::handlePacket);
+
+    InetSocketAddress address =
+        new InetSocketAddress(socket.localAddress().host(), socket.localAddress().port());
+    addressFuture.complete(address);
+  }
+
+  @Override
+  protected CompletableFuture<Void> sendOutgoingPacket(
+      final DiscoveryPeer peer, final Packet packet) {
+    CompletableFuture<Void> result = new CompletableFuture<>();
+    socket.send(
+        packet.encode(),
+        peer.getEndpoint().getUdpPort(),
+        peer.getEndpoint().getHost(),
+        ar -> {
+          if (ar.failed()) {
+            result.completeExceptionally(ar.cause());
+          } else {
+            result.complete(null);
+          }
+        });
+    return result;
+  }
+
+  @Override
+  public CompletableFuture<?> stop() {
+    if (socket == null) {
+      return CompletableFuture.completedFuture(null);
+    }
+
+    final CompletableFuture<?> completion = new CompletableFuture<>();
+    socket.close(
+        ar -> {
+          if (ar.succeeded()) {
+            controller.ifPresent(PeerDiscoveryController::stop);
+            socket = null;
+            completion.complete(null);
+          } else {
+            completion.completeExceptionally(ar.cause());
+          }
+        });
+    return completion;
+  }
+
+  /**
+   * For uncontrolled exceptions occurring in the packet handlers.
+   *
+   * @param exception the exception that was raised
+   */
+  private void handleException(final Throwable exception) {
+    if (exception instanceof IOException) {
+      LOG.debug("Packet handler exception", exception);
+    } else {
+      LOG.error("Packet handler exception", exception);
+    }
+  }
+
+  /**
+   * The UDP packet handler. This is the entrypoint for all received datagrams.
+   *
+   * @param datagram the received datagram.
+   */
+  private void handlePacket(final DatagramPacket datagram) {
+    try {
+      final int length = datagram.data().length();
+      Preconditions.checkGuard(
+          validatePacketSize(length),
+          PeerDiscoveryPacketDecodingException::new,
+          "Packet too large. Actual size (bytes): %s",
+          length);
+
+      // We allow exceptions to bubble up, as they'll be picked up by the exception handler.
+      final Packet packet = Packet.decode(datagram.data());
+      // Acquire the senders coordinates to build a Peer representation from them.
+      final String host = datagram.sender().host();
+      final int port = datagram.sender().port();
+      final Endpoint endpoint = new Endpoint(host, port, OptionalInt.empty());
+      handleIncomingPacket(endpoint, packet);
+    } catch (final PeerDiscoveryPacketDecodingException e) {
+      LOG.debug("Discarding invalid peer discovery packet", e);
+    } catch (final Throwable t) {
+      LOG.error("Encountered error while handling packet", t);
+    }
+  }
+}

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/OutboundMessageHandler.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
+
+import tech.pegasys.pantheon.ethereum.p2p.discovery.DiscoveryPeer;
+
+@FunctionalInterface
+public interface OutboundMessageHandler {
+  public static OutboundMessageHandler NOOP = (peer, packet) -> {};
+
+  void send(final DiscoveryPeer toPeer, final Packet packet);
+}

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/PeerDiscoveryController.java

@@ -12,17 +12,16 @@
  */
 package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
 
-import static com.google.common.base.Preconditions.checkNotNull;
 import static java.util.Collections.emptyList;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerTable.AddResult.Outcome;
 
+import tech.pegasys.pantheon.crypto.SECP256K1;
+import tech.pegasys.pantheon.crypto.SECP256K1.KeyPair;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.DiscoveryPeer;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryAgent;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryEvent;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryEvent.PeerBondedEvent;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryEvent.PeerDroppedEvent;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryStatus;
 import tech.pegasys.pantheon.ethereum.p2p.peers.Peer;
 import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
@@ -42,7 +41,6 @@ import java.util.function.Consumer;
 import java.util.function.Predicate;
 
 import com.google.common.annotations.VisibleForTesting;
-import io.vertx.core.Vertx;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -94,7 +92,7 @@ public class PeerDiscoveryController {
 
   private static final Logger LOG = LogManager.getLogger();
   private static final long REFRESH_CHECK_INTERVAL_MILLIS = MILLISECONDS.convert(30, SECONDS);
-  private final Vertx vertx;
+  protected final TimerUtil timerUtil;
   private final PeerTable peerTable;
 
   private final Collection<DiscoveryPeer> bootstrapNodes;
@@ -105,7 +103,10 @@ public class PeerDiscoveryController {
 
   private final AtomicBoolean started = new AtomicBoolean(false);
 
-  private final PeerDiscoveryAgent agent;
+  private final SECP256K1.KeyPair keypair;
+  // The peer representation of this node
+  private final DiscoveryPeer localPeer;
+  private final OutboundMessageHandler outboundMessageHandler;
   private final PeerBlacklist peerBlacklist;
   private final NodeWhitelistController nodeWhitelist;
 
@@ -120,28 +121,31 @@ public class PeerDiscoveryController {
   private OptionalLong tableRefreshTimerId = OptionalLong.empty();
 
   // Observers for "peer bonded" discovery events.
-  private final Subscribers<Consumer<PeerBondedEvent>> peerBondedObservers = new Subscribers<>();
+  private final Subscribers<Consumer<PeerBondedEvent>> peerBondedObservers;
-
-  // Observers for "peer dropped" discovery events.
-  private final Subscribers<Consumer<PeerDroppedEvent>> peerDroppedObservers = new Subscribers<>();
 
   public PeerDiscoveryController(
-      final Vertx vertx,
+      final KeyPair keypair,
-      final PeerDiscoveryAgent agent,
+      final DiscoveryPeer localPeer,
       final PeerTable peerTable,
       final Collection<DiscoveryPeer> bootstrapNodes,
+      final OutboundMessageHandler outboundMessageHandler,
+      final TimerUtil timerUtil,
       final long tableRefreshIntervalMs,
       final PeerRequirement peerRequirement,
       final PeerBlacklist peerBlacklist,
-      final NodeWhitelistController nodeWhitelist) {
+      final NodeWhitelistController nodeWhitelist,
-    this.vertx = vertx;
+      final Subscribers<Consumer<PeerBondedEvent>> peerBondedObservers) {
-    this.agent = agent;
+    this.timerUtil = timerUtil;
+    this.keypair = keypair;
+    this.localPeer = localPeer;
     this.bootstrapNodes = bootstrapNodes;
     this.peerTable = peerTable;
     this.tableRefreshIntervalMs = tableRefreshIntervalMs;
     this.peerRequirement = peerRequirement;
     this.peerBlacklist = peerBlacklist;
     this.nodeWhitelist = nodeWhitelist;
+    this.outboundMessageHandler = outboundMessageHandler;
+    this.peerBondedObservers = peerBondedObservers;
   }
 
   public CompletableFuture<?> start() {
@@ -156,9 +160,9 @@ public class PeerDiscoveryController {
         .forEach(node -> bond(node, true));
 
     final long timerId =
-        vertx.setPeriodic(
+        timerUtil.setPeriodic(
             Math.min(REFRESH_CHECK_INTERVAL_MILLIS, tableRefreshIntervalMs),
-            (l) -> refreshTableIfRequired());
+            () -> refreshTableIfRequired());
     tableRefreshTimerId = OptionalLong.of(timerId);
 
     return CompletableFuture.completedFuture(null);
@@ -169,7 +173,7 @@ public class PeerDiscoveryController {
       return CompletableFuture.completedFuture(null);
     }
 
-    tableRefreshTimerId.ifPresent(vertx::cancelTimer);
+    tableRefreshTimerId.ifPresent(timerUtil::cancelTimer);
     tableRefreshTimerId = OptionalLong.empty();
     inflightInteractions.values().forEach(PeerInteractionState::cancelTimers);
     inflightInteractions.clear();
@@ -196,7 +200,7 @@ public class PeerDiscoveryController {
         packet);
 
     // Message from self. This should not happen.
-    if (sender.getId().equals(agent.getAdvertisedPeer().getId())) {
+    if (sender.getId().equals(localPeer.getId())) {
       return;
     }
 
@@ -230,7 +234,7 @@ public class PeerDiscoveryController {
 
                     // If this was a bootstrap peer, let's ask it for nodes near to us.
                     if (interaction.isBootstrap()) {
-                      findNodes(peer, agent.getAdvertisedPeer().getId());
+                      findNodes(peer, localPeer.getId());
                     }
                   });
           break;
@@ -247,9 +251,12 @@ public class PeerDiscoveryController {
                           .orElse(emptyList());
 
                   for (final DiscoveryPeer neighbor : neighbors) {
+                    // If the peer is not whitelisted, is blacklisted, is already known, or
+                    // represents this node, skip bonding
                     if (!nodeWhitelist.isPermitted(neighbor)
                         || peerBlacklist.contains(neighbor)
-                        || peerTable.get(neighbor).isPresent()) {
+                        || peerTable.get(neighbor).isPresent()
+                        || neighbor.getId().equals(localPeer.getId())) {
                       continue;
                     }
                     bond(neighbor, false);
@@ -315,7 +322,7 @@ public class PeerDiscoveryController {
 
   private void refreshTableIfRequired() {
     final long now = System.currentTimeMillis();
-    if (lastRefreshTime + tableRefreshIntervalMs < now) {
+    if (lastRefreshTime + tableRefreshIntervalMs <= now) {
       LOG.info("Peer table refresh triggered by timer expiry");
       refreshTable();
     } else if (!peerRequirement.hasSufficientPeers()) {
@@ -348,10 +355,10 @@ public class PeerDiscoveryController {
     final Consumer<PeerInteractionState> action =
         interaction -> {
           final PingPacketData data =
-              PingPacketData.create(agent.getAdvertisedPeer().getEndpoint(), peer.getEndpoint());
+              PingPacketData.create(localPeer.getEndpoint(), peer.getEndpoint());
-          final Packet sentPacket = agent.sendPacket(peer, PacketType.PING, data);
+          final Packet pingPacket = createPacket(PacketType.PING, data);
 
-          final BytesValue pingHash = sentPacket.getHash();
+          final BytesValue pingHash = pingPacket.getHash();
           // Update the matching filter to only accept the PONG if it echoes the hash of our PING.
           final Predicate<Packet> newFilter =
               packet ->
@@ -360,6 +367,8 @@ public class PeerDiscoveryController {
                       .map(pong -> pong.getPingHash().equals(pingHash))
                       .orElse(false);
           interaction.updateFilter(newFilter);
+
+          sendPacket(peer, pingPacket);
         };
 
     // The filter condition will be updated as soon as the action is performed.
@@ -368,6 +377,20 @@ public class PeerDiscoveryController {
     dispatchInteraction(peer, ping);
   }
 
+  private void sendPacket(final DiscoveryPeer peer, final PacketType type, final PacketData data) {
+    Packet packet = createPacket(type, data);
+    outboundMessageHandler.send(peer, packet);
+  }
+
+  private void sendPacket(final DiscoveryPeer peer, final Packet packet) {
+    outboundMessageHandler.send(peer, packet);
+  }
+
+  @VisibleForTesting
+  Packet createPacket(final PacketType type, final PacketData data) {
+    return Packet.create(type, data, keypair);
+  }
+
   /**
    * Sends a FIND_NEIGHBORS message to a {@link DiscoveryPeer}, in search of a target value.
    *
@@ -378,7 +401,7 @@ public class PeerDiscoveryController {
     final Consumer<PeerInteractionState> action =
         (interaction) -> {
           final FindNeighborsPacketData data = FindNeighborsPacketData.create(target);
-          agent.sendPacket(peer, PacketType.FIND_NEIGHBORS, data);
+          sendPacket(peer, PacketType.FIND_NEIGHBORS, data);
         };
     final PeerInteractionState interaction =
         new PeerInteractionState(action, PacketType.NEIGHBORS, packet -> true, true, false);
@@ -405,7 +428,7 @@ public class PeerDiscoveryController {
   private void respondToPing(
       final PingPacketData packetData, final BytesValue pingHash, final DiscoveryPeer sender) {
     final PongPacketData data = PongPacketData.create(packetData.getFrom(), pingHash);
-    agent.sendPacket(sender, PacketType.PONG, data);
+    sendPacket(sender, PacketType.PONG, data);
   }
 
   private void respondToFindNeighbors(
@@ -414,22 +437,13 @@ public class PeerDiscoveryController {
     // peers they can fit in a 1280-byte payload.
     final List<DiscoveryPeer> peers = peerTable.nearestPeers(packetData.getTarget(), 16);
     final PacketData data = NeighborsPacketData.create(peers);
-    agent.sendPacket(sender, PacketType.NEIGHBORS, data);
+    sendPacket(sender, PacketType.NEIGHBORS, data);
   }
 
-  // Dispatches an event to a set of observers. Since we have no control over observer logic, we
+  // Dispatches an event to a set of observers.
-  // take
-  // precautions and we assume they are of blocking nature to protect our event loop.
   private <T extends PeerDiscoveryEvent> void dispatchEvent(
       final Subscribers<Consumer<T>> observers, final T event) {
-    observers.forEach(
+    observers.forEach(observer -> observer.accept(event));
-        observer ->
-            vertx.executeBlocking(
-                future -> {
-                  observer.accept(event);
-                  future.complete();
-                },
-                x -> {}));
   }
 
   /**
@@ -446,63 +460,6 @@ public class PeerDiscoveryController {
     this.retryDelayFunction = retryDelayFunction;
   }
 
-  /**
-   * Adds an observer that will get called when a new peer is bonded with and added to the peer
-   * table.
-   *
-   * <p><i>No guarantees are made about the order in which observers are invoked.</i>
-   *
-   * @param observer The observer to call.
-   * @return A unique ID identifying this observer, to that it can be removed later.
-   */
-  public long observePeerBondedEvents(final Consumer<PeerBondedEvent> observer) {
-    checkNotNull(observer);
-    return peerBondedObservers.subscribe(observer);
-  }
-
-  /**
-   * Adds an observer that will get called when a new peer is dropped from the peer table.
-   *
-   * <p><i>No guarantees are made about the order in which observers are invoked.</i>
-   *
-   * @param observer The observer to call.
-   * @return A unique ID identifying this observer, to that it can be removed later.
-   */
-  public long observePeerDroppedEvents(final Consumer<PeerDroppedEvent> observer) {
-    checkNotNull(observer);
-    return peerDroppedObservers.subscribe(observer);
-  }
-
-  /**
-   * Removes an previously added peer bonded observer.
-   *
-   * @param observerId The unique ID identifying the observer to remove.
-   * @return Whether the observer was located and removed.
-   */
-  public boolean removePeerBondedObserver(final long observerId) {
-    return peerBondedObservers.unsubscribe(observerId);
-  }
-
-  /**
-   * Removes an previously added peer dropped observer.
-   *
-   * @param observerId The unique ID identifying the observer to remove.
-   * @return Whether the observer was located and removed.
-   */
-  public boolean removePeerDroppedObserver(final long observerId) {
-    return peerDroppedObservers.unsubscribe(observerId);
-  }
-
-  /**
-   * Returns the count of observers that are registered on this controller.
-   *
-   * @return The observer count.
-   */
-  @VisibleForTesting
-  public int observerCount() {
-    return peerBondedObservers.getSubscriberCount() + peerDroppedObservers.getSubscriberCount();
-  }
-
   /** Holds the state machine data for a peer interaction. */
   private class PeerInteractionState implements Predicate<Packet> {
     /**
@@ -558,13 +515,13 @@ public class PeerDiscoveryController {
       action.accept(this);
       if (retryable) {
         final long newTimeout = retryDelayFunction.apply(lastTimeout);
-        timerId = OptionalLong.of(vertx.setTimer(newTimeout, id -> execute(newTimeout)));
+        timerId = OptionalLong.of(timerUtil.setTimer(newTimeout, () -> execute(newTimeout)));
       }
     }
 
     /** Cancels any timers associated with this entry. */
     void cancelTimers() {
-      timerId.ifPresent(vertx::cancelTimer);
+      timerId.ifPresent(timerUtil::cancelTimer);
     }
   }
 }

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/PeerRequirement.java

@@ -14,6 +14,7 @@ package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
 
 import java.util.Collection;
 
+@FunctionalInterface
 public interface PeerRequirement {
 
   boolean hasSufficientPeers();

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/TimerUtil.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
+
+public interface TimerUtil {
+  long setPeriodic(long delay, TimerHandler handler);
+
+  long setTimer(long delay, TimerHandler handler);
+
+  void cancelTimer(long timerId);
+
+  @FunctionalInterface
+  interface TimerHandler {
+    void handle();
+  }
+}

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/VertxTimerUtil.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
+
+import io.vertx.core.Vertx;
+
+public class VertxTimerUtil implements TimerUtil {
+
+  private final Vertx vertx;
+
+  public VertxTimerUtil(final Vertx vertx) {
+    this.vertx = vertx;
+  }
+
+  @Override
+  public long setPeriodic(final long delay, final TimerHandler handler) {
+    return vertx.setPeriodic(delay, (l) -> handler.handle());
+  }
+
+  @Override
+  public long setTimer(final long delay, final TimerHandler handler) {
+    return vertx.setTimer(delay, (l) -> handler.handle());
+  }
+
+  @Override
+  public void cancelTimer(final long timerId) {
+    vertx.cancelTimer(timerId);
+  }
+}

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/netty/NettyP2PNetwork.java

@@ -22,6 +22,7 @@ import tech.pegasys.pantheon.ethereum.p2p.api.PeerConnection;
 import tech.pegasys.pantheon.ethereum.p2p.config.NetworkingConfiguration;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.DiscoveryPeer;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryAgent;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.VertxPeerDiscoveryAgent;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerRequirement;
 import tech.pegasys.pantheon.ethereum.p2p.peers.Endpoint;
 import tech.pegasys.pantheon.ethereum.p2p.peers.Peer;
@@ -50,6 +51,7 @@ import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.netty.bootstrap.Bootstrap;
 import io.netty.bootstrap.ServerBootstrap;
 import io.netty.channel.ChannelFuture;
@@ -174,7 +176,7 @@ public final class NettyP2PNetwork implements P2PNetwork {
     this.peerBlacklist = peerBlacklist;
     this.nodeWhitelistController = nodeWhitelistController;
     peerDiscoveryAgent =
-        new PeerDiscoveryAgent(
+        new VertxPeerDiscoveryAgent(
             vertx,
             keyPair,
             config.getDiscovery(),
@@ -380,7 +382,7 @@ public final class NettyP2PNetwork implements P2PNetwork {
   @Override
   public void run() {
     try {
-      peerDiscoveryAgent.start(ourPeerInfo.getPort()).join();
+      peerDiscoveryAgent.start().join();
       final long observerId =
           peerDiscoveryAgent.observePeerBondedEvents(
               peerBondedEvent -> {
@@ -425,6 +427,7 @@ public final class NettyP2PNetwork implements P2PNetwork {
     stop();
   }
 
+  @VisibleForTesting
   public Collection<DiscoveryPeer> getDiscoveryPeers() {
     return peerDiscoveryAgent.getPeers();
   }
@@ -435,7 +438,7 @@ public final class NettyP2PNetwork implements P2PNetwork {
   }
 
   @Override
-  public PeerInfo getSelf() {
+  public PeerInfo getLocalPeerInfo() {
     return ourPeerInfo;
   }
 

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/NettyP2PNetworkTest.java

@@ -93,7 +93,7 @@ public final class NettyP2PNetworkTest {
                 new NoOpMetricsSystem(),
                 new NodeWhitelistController(PermissioningConfiguration.createDefault()))) {
 
-      final int listenPort = listener.getSelf().getPort();
+      final int listenPort = listener.getLocalPeerInfo().getPort();
       listener.run();
       connector.run();
       final BytesValue listenId = listenKp.getPublicKey().getEncodedBytes();
@@ -146,7 +146,7 @@ public final class NettyP2PNetworkTest {
                 new PeerBlacklist(),
                 new NoOpMetricsSystem(),
                 new NodeWhitelistController(PermissioningConfiguration.createDefault()))) {
-      final int listenPort = listener.getSelf().getPort();
+      final int listenPort = listener.getLocalPeerInfo().getPort();
       listener.run();
       connector.run();
       final BytesValue listenId = listenKp.getPublicKey().getEncodedBytes();
@@ -229,7 +229,7 @@ public final class NettyP2PNetworkTest {
                 new NoOpMetricsSystem(),
                 new NodeWhitelistController(PermissioningConfiguration.createDefault()))) {
 
-      final int listenPort = listener.getSelf().getPort();
+      final int listenPort = listener.getLocalPeerInfo().getPort();
       // Setup listener and first connection
       listener.run();
       connector1.run();
@@ -296,7 +296,7 @@ public final class NettyP2PNetworkTest {
                 new PeerBlacklist(),
                 new NoOpMetricsSystem(),
                 new NodeWhitelistController(PermissioningConfiguration.createDefault()))) {
-      final int listenPort = listener.getSelf().getPort();
+      final int listenPort = listener.getLocalPeerInfo().getPort();
       listener.run();
       connector.run();
       final BytesValue listenId = listenKp.getPublicKey().getEncodedBytes();
@@ -351,8 +351,8 @@ public final class NettyP2PNetworkTest {
                 remoteBlacklist,
                 new NoOpMetricsSystem(),
                 new NodeWhitelistController(PermissioningConfiguration.createDefault()))) {
-      final int localListenPort = localNetwork.getSelf().getPort();
+      final int localListenPort = localNetwork.getLocalPeerInfo().getPort();
-      final int remoteListenPort = remoteNetwork.getSelf().getPort();
+      final int remoteListenPort = remoteNetwork.getLocalPeerInfo().getPort();
       final Peer localPeer =
           new DefaultPeer(
               localId,

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/AbstractPeerDiscoveryTest.java

@@ -1,293 +0,0 @@
-/*
- * Copyright 2018 ConsenSys AG.
- *
- * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
- * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations under the License.
- */
-package tech.pegasys.pantheon.ethereum.p2p.discovery;
-
-import static io.vertx.core.Vertx.vertx;
-
-import tech.pegasys.pantheon.crypto.SECP256K1;
-import tech.pegasys.pantheon.ethereum.p2p.config.DiscoveryConfiguration;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketType;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PingPacketData;
-import tech.pegasys.pantheon.ethereum.p2p.peers.Endpoint;
-import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
-import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
-import tech.pegasys.pantheon.ethereum.permissioning.PermissioningConfiguration;
-import tech.pegasys.pantheon.util.bytes.BytesValue;
-
-import java.util.List;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
-
-import io.vertx.core.Vertx;
-import io.vertx.core.datagram.DatagramSocket;
-import junit.framework.AssertionFailedError;
-import org.junit.After;
-
-/**
- * A test class you can extend to acquire the ability to easily start discovery agents with a
- * generated Peer and keypair, as well as test sockets to communicate with those discovery agents.
- *
- * <p>Call {@link #startDiscoveryAgent(List)} and variants to start one or more discovery agents, or
- * {@link #startTestSocket()} or variants to start one or more test sockets. The lifecycle of those
- * objects is managed automatically for you via @Before and @After hooks, so you don't need to worry
- * about starting and stopping.
- */
-public abstract class AbstractPeerDiscoveryTest {
-  private static final String LOOPBACK_IP_ADDR = "127.0.0.1";
-  private static final int TEST_SOCKET_START_TIMEOUT_SECS = 5;
-  private static final int BROADCAST_TCP_PORT = 12356;
-  private final Vertx vertx = vertx();
-
-  List<DiscoveryTestSocket> discoveryTestSockets = new CopyOnWriteArrayList<>();
-  List<PeerDiscoveryAgent> agents = new CopyOnWriteArrayList<>();
-
-  @After
-  public void stopServices() {
-    // Close all sockets, will bubble up exceptions.
-    final CompletableFuture<?>[] completions =
-        discoveryTestSockets
-            .stream()
-            .filter(p -> p.getSocket() != null)
-            .map(
-                p -> {
-                  final CompletableFuture<?> completion = new CompletableFuture<>();
-                  p.getSocket()
-                      .close(
-                          ar -> {
-                            if (ar.succeeded()) {
-                              completion.complete(null);
-                            } else {
-                              completion.completeExceptionally(ar.cause());
-                            }
-                          });
-                  return completion;
-                })
-            .toArray(CompletableFuture<?>[]::new);
-    try {
-      CompletableFuture.allOf(completions).join();
-    } finally {
-      agents.forEach(PeerDiscoveryAgent::stop);
-      vertx.close();
-    }
-  }
-
-  /**
-   * Starts multiple discovery agents with the provided boostrap peers.
-   *
-   * @param count the number of agents to start
-   * @param bootstrapPeers the list of bootstrap peers
-   * @return a list of discovery agents.
-   */
-  protected List<PeerDiscoveryAgent> startDiscoveryAgents(
-      final int count, final List<DiscoveryPeer> bootstrapPeers) {
-    return Stream.generate(() -> startDiscoveryAgent(bootstrapPeers))
-        .limit(count)
-        .collect(Collectors.toList());
-  }
-
-  /**
-   * Start a single discovery agent with the provided bootstrap peers.
-   *
-   * @param bootstrapPeers the list of bootstrap peers
-   * @return a list of discovery agents.
-   */
-  protected PeerDiscoveryAgent startDiscoveryAgent(final List<DiscoveryPeer> bootstrapPeers) {
-    return startDiscoveryAgent(bootstrapPeers, new PeerBlacklist());
-  }
-
-  /**
-   * Start a single discovery agent with the provided bootstrap peers.
-   *
-   * @param bootstrapPeers the list of bootstrap peers
-   * @param blacklist the peer blacklist
-   * @return a list of discovery agents.
-   */
-  protected PeerDiscoveryAgent startDiscoveryAgent(
-      final List<DiscoveryPeer> bootstrapPeers, final PeerBlacklist blacklist) {
-    final DiscoveryConfiguration config = new DiscoveryConfiguration();
-    config.setBootstrapPeers(bootstrapPeers);
-    config.setBindPort(0);
-
-    return startDiscoveryAgent(config, blacklist);
-  }
-
-  protected PeerDiscoveryAgent startDiscoveryAgent(
-      final DiscoveryConfiguration config, final PeerBlacklist blacklist) {
-    final PeerDiscoveryAgent agent =
-        new PeerDiscoveryAgent(
-            vertx,
-            SECP256K1.KeyPair.generate(),
-            config,
-            () -> true,
-            blacklist,
-            new NodeWhitelistController(PermissioningConfiguration.createDefault()));
-    try {
-      agent.start(BROADCAST_TCP_PORT).get(5, TimeUnit.SECONDS);
-    } catch (final Exception ex) {
-      throw new AssertionError("Could not initialize discovery agent", ex);
-    }
-    agents.add(agent);
-    return agent;
-  }
-
-  /**
-   * Start multiple test sockets.
-   *
-   * <p>A test socket allows you to send messages to a discovery agent, as well as to react to
-   * received messages. A test socket encapsulates: (1) a {@link DiscoveryPeer} and its {@link
-   * tech.pegasys.pantheon.crypto.SECP256K1.KeyPair}, (2) an {@link ArrayBlockingQueue} where
-   * received messages are placed automatically, and (3) the socket itself.
-   *
-   * @param count the number of test sockets to start.
-   * @return the test sockets.
-   */
-  protected List<DiscoveryTestSocket> startTestSockets(final int count) {
-    return Stream.generate(this::startTestSocket).limit(count).collect(Collectors.toList());
-  }
-
-  /**
-   * Starts a single test socket.
-   *
-   * @return the test socket
-   */
-  protected DiscoveryTestSocket startTestSocket() {
-    final ArrayBlockingQueue<Packet> queue = new ArrayBlockingQueue<>(100);
-
-    final SECP256K1.KeyPair keyPair = SECP256K1.KeyPair.generate();
-    final BytesValue peerId = keyPair.getPublicKey().getEncodedBytes();
-
-    final CompletableFuture<DiscoveryTestSocket> result = new CompletableFuture<>();
-    // Test packet handler which feeds the received packet into a Future we later consume from.
-    vertx
-        .createDatagramSocket()
-        .listen(
-            0,
-            LOOPBACK_IP_ADDR,
-            ar -> {
-              if (!ar.succeeded()) {
-                result.completeExceptionally(ar.cause());
-                return;
-              }
-
-              final DatagramSocket socket = ar.result();
-              socket.handler(p -> queue.add(Packet.decode(p.data())));
-              final DiscoveryPeer peer =
-                  new DiscoveryPeer(
-                      peerId,
-                      LOOPBACK_IP_ADDR,
-                      socket.localAddress().port(),
-                      socket.localAddress().port());
-              final DiscoveryTestSocket discoveryTestSocket =
-                  new DiscoveryTestSocket(peer, keyPair, queue, socket);
-              result.complete(discoveryTestSocket);
-            });
-
-    final DiscoveryTestSocket discoveryTestSocket;
-    try {
-      discoveryTestSocket = result.get(TEST_SOCKET_START_TIMEOUT_SECS, TimeUnit.SECONDS);
-    } catch (final Exception ex) {
-      throw new AssertionError("Could not initialize test peer", ex);
-    }
-    discoveryTestSockets.add(discoveryTestSocket);
-    return discoveryTestSocket;
-  }
-
-  protected void bondViaIncomingPing(
-      final PeerDiscoveryAgent agent, final DiscoveryTestSocket peerSocket)
-      throws InterruptedException {
-    final DiscoveryPeer peer = peerSocket.getPeer();
-
-    final PingPacketData ping =
-        PingPacketData.create(peer.getEndpoint(), agent.getAdvertisedPeer().getEndpoint());
-    final Packet pingPacket = Packet.create(PacketType.PING, ping, peerSocket.getKeyPair());
-    peerSocket.sendToAgent(agent, pingPacket);
-
-    // Wait for returned pong packet to finish bonding
-    peerSocket.getIncomingPackets().poll(10, TimeUnit.SECONDS);
-  }
-
-  /**
-   * Encapsulates a test socket representing a Peer, with an associated queue where all incoming
-   * packets are placed.
-   */
-  protected static class DiscoveryTestSocket {
-    private final DiscoveryPeer peer;
-    private final SECP256K1.KeyPair keyPair;
-    private final ArrayBlockingQueue<Packet> queue;
-    private final DatagramSocket socket;
-
-    public DiscoveryTestSocket(
-        final DiscoveryPeer peer,
-        final SECP256K1.KeyPair keyPair,
-        final ArrayBlockingQueue<Packet> queue,
-        final DatagramSocket socket) {
-      this.peer = peer;
-      this.keyPair = keyPair;
-      this.queue = queue;
-      this.socket = socket;
-    }
-
-    public DiscoveryPeer getPeer() {
-      return peer;
-    }
-
-    public ArrayBlockingQueue<Packet> getIncomingPackets() {
-      return queue;
-    }
-
-    public DatagramSocket getSocket() {
-      return socket;
-    }
-
-    public SECP256K1.KeyPair getKeyPair() {
-      return keyPair;
-    }
-
-    /**
-     * Sends a message to an agent.
-     *
-     * @param agent the recipient
-     * @param packet the packet to send
-     */
-    public void sendToAgent(final PeerDiscoveryAgent agent, final Packet packet) {
-      final Endpoint endpoint = agent.getAdvertisedPeer().getEndpoint();
-      socket.send(packet.encode(), endpoint.getUdpPort(), endpoint.getHost(), ar -> {});
-    }
-
-    /**
-     * Retrieves the head of the queue, compulsorily. If no message exists, or no message appears in
-     * 5 seconds, it throws an assertion error.
-     *
-     * @return the head of the queue
-     */
-    public Packet compulsoryPoll() {
-      final Packet packet;
-      try {
-        packet = queue.poll(5, TimeUnit.SECONDS);
-      } catch (final Exception e) {
-        throw new RuntimeException(e);
-      }
-
-      if (packet == null) {
-        throw new AssertionFailedError(
-            "Expected a message in the test peer queue, but found none.");
-      }
-      return packet;
-    }
-  }
-}

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryAgentTest.java

@@ -13,152 +13,116 @@
 package tech.pegasys.pantheon.ethereum.p2p.discovery;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.awaitility.Awaitility.await;
 
-import tech.pegasys.pantheon.crypto.SECP256K1;
 import tech.pegasys.pantheon.ethereum.p2p.api.MessageData;
 import tech.pegasys.pantheon.ethereum.p2p.api.PeerConnection;
-import tech.pegasys.pantheon.ethereum.p2p.config.DiscoveryConfiguration;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper.AgentBuilder;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.FindNeighborsPacketData;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent.IncomingPacket;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.NeighborsPacketData;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketType;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PingPacketData;
-import tech.pegasys.pantheon.ethereum.p2p.peers.DefaultPeer;
 import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
-import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
 import tech.pegasys.pantheon.ethereum.p2p.wire.Capability;
 import tech.pegasys.pantheon.ethereum.p2p.wire.PeerInfo;
 import tech.pegasys.pantheon.ethereum.p2p.wire.messages.DisconnectMessage.DisconnectReason;
-import tech.pegasys.pantheon.ethereum.permissioning.PermissioningConfiguration;
 import tech.pegasys.pantheon.util.bytes.BytesValue;
 
 import java.net.SocketAddress;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
-import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
-import com.google.common.collect.Lists;
-import io.vertx.core.Vertx;
-import org.junit.Ignore;
 import org.junit.Test;
 
-public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
+public class PeerDiscoveryAgentTest {
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
-  public void neighborsPacketFromUnbondedPeerIsDropped() throws Exception {
+  public void neighborsPacketFromUnbondedPeerIsDropped() {
     // Start an agent with no bootstrap peers.
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList());
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(Collections.emptyList());
     assertThat(agent.getPeers()).isEmpty();
 
-    // Start a test peer and send a PING packet to the agent under test.
+    // Start a test peer
-    final DiscoveryTestSocket discoveryTestSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherNode = helper.startDiscoveryAgent();
-
-    // Peer is unbonded, as it has not replied with a PONG.
 
     // Generate an out-of-band NEIGHBORS message.
-    final DiscoveryPeer[] peers =
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(5);
-        PeerDiscoveryTestHelper.generatePeers(PeerDiscoveryTestHelper.generateKeyPairs(5));
+    final NeighborsPacketData data = NeighborsPacketData.create(peers);
-    final NeighborsPacketData data = NeighborsPacketData.create(Arrays.asList(peers));
+    final Packet packet = Packet.create(PacketType.NEIGHBORS, data, otherNode.getKeyPair());
-    final Packet packet =
+    helper.sendMessageBetweenAgents(otherNode, agent, packet);
-        Packet.create(PacketType.NEIGHBORS, data, discoveryTestSocket.getKeyPair());
+
-    discoveryTestSocket.sendToAgent(agent, packet);
-
-    TimeUnit.SECONDS.sleep(1);
     assertThat(agent.getPeers()).isEmpty();
   }
 
   @Test
-  @Ignore("This test is failing intermittently - disabling while we investigate")
   public void neighborsPacketLimited() {
     // Start 20 agents with no bootstrap peers.
-    final List<PeerDiscoveryAgent> agents = startDiscoveryAgents(20, Collections.emptyList());
+    final List<MockPeerDiscoveryAgent> otherAgents =
-    final List<DiscoveryPeer> agentPeers =
+        helper.startDiscoveryAgents(20, Collections.emptyList());
-        agents.stream().map(PeerDiscoveryAgent::getAdvertisedPeer).collect(Collectors.toList());
+    final List<DiscoveryPeer> otherPeers =
-
+        otherAgents
-    // Start another bootstrap peer pointing to those 20 agents.
+            .stream()
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(agentPeers);
+            .map(MockPeerDiscoveryAgent::getAdvertisedPeer)
-    await()
+            .collect(Collectors.toList());
-        .atMost(10, TimeUnit.SECONDS)
+
-        .untilAsserted(
+    // Start another peer pointing to those 20 agents.
-            () -> {
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(otherPeers);
     assertThat(agent.getPeers()).hasSize(20);
-              assertThat(agent.getPeers())
+    assertThat(agent.getPeers()).allMatch(p -> p.getStatus() == PeerDiscoveryStatus.BONDED);
-                  .allMatch(p -> p.getStatus() == PeerDiscoveryStatus.BONDED);
-            });
 
-    // Send a PING so we can exchange messages with the latter agent.
+    // Use additional agent to exchange messages with agent
-    final DiscoveryTestSocket testSocket = startTestSocket();
+    final MockPeerDiscoveryAgent testAgent = helper.startDiscoveryAgent();
-    Packet packet =
-        Packet.create(
-            PacketType.PING,
-            PingPacketData.create(
-                testSocket.getPeer().getEndpoint(), testSocket.getPeer().getEndpoint()),
-            testSocket.getKeyPair());
-    testSocket.sendToAgent(agent, packet);
 
-    // Wait until PONG is received.
+    // Send a PING so we can exchange messages with the latter agent.
-    final Packet pong = testSocket.compulsoryPoll();
+    Packet packet = helper.createPingPacket(testAgent, agent);
-    assertThat(pong.getType()).isEqualTo(PacketType.PONG);
+    helper.sendMessageBetweenAgents(testAgent, agent, packet);
 
     // Send a FIND_NEIGHBORS message.
     packet =
         Packet.create(
             PacketType.FIND_NEIGHBORS,
-            FindNeighborsPacketData.create(agents.get(0).getAdvertisedPeer().getId()),
+            FindNeighborsPacketData.create(otherAgents.get(0).getAdvertisedPeer().getId()),
-            testSocket.getKeyPair());
+            testAgent.getKeyPair());
-    testSocket.sendToAgent(agent, packet);
+    helper.sendMessageBetweenAgents(testAgent, agent, packet);
-
+
-    // Wait until NEIGHBORS is received.
+    // Check response packet
-    packet = testSocket.compulsoryPoll();
+    List<IncomingPacket> incomingPackets =
-    assertThat(packet.getType()).isEqualTo(PacketType.NEIGHBORS);
+        testAgent
+            .getIncomingPackets()
+            .stream()
+            .filter(p -> p.packet.getType().equals(PacketType.NEIGHBORS))
+            .collect(Collectors.toList());
+    assertThat(incomingPackets.size()).isEqualTo(1);
+    IncomingPacket neighborsPacket = incomingPackets.get(0);
+    assertThat(neighborsPacket.fromAgent).isEqualTo(agent);
 
     // Assert that we only received 16 items.
-    final NeighborsPacketData neighbors = packet.getPacketData(NeighborsPacketData.class).get();
+    final NeighborsPacketData neighbors =
+        neighborsPacket.packet.getPacketData(NeighborsPacketData.class).get();
     assertThat(neighbors).isNotNull();
     assertThat(neighbors.getNodes()).hasSize(16);
 
     // Assert that after removing those 16 items we're left with either 4 or 5.
     // If we are left with 5, the test peer was returned as an item, assert that this is the case.
-    agentPeers.removeAll(neighbors.getNodes());
+    otherPeers.removeAll(neighbors.getNodes());
-    assertThat(agentPeers.size()).isBetween(4, 5);
+    assertThat(otherPeers.size()).isBetween(4, 5);
-    if (agentPeers.size() == 5) {
+    if (otherPeers.size() == 5) {
-      assertThat(neighbors.getNodes()).contains(testSocket.getPeer());
+      assertThat(neighbors.getNodes()).contains(testAgent.getAdvertisedPeer());
     }
   }
 
   @Test
   public void shouldEvictPeerOnDisconnect() {
-    final Vertx vertx = Vertx.vertx();
+    final MockPeerDiscoveryAgent peerDiscoveryAgent1 = helper.startDiscoveryAgent();
-
+    peerDiscoveryAgent1.start().join();
-    final SECP256K1.KeyPair keyPair1 = SECP256K1.KeyPair.generate();
+    final DiscoveryPeer peer = peerDiscoveryAgent1.getAdvertisedPeer();
-    final PeerDiscoveryAgent peerDiscoveryAgent1 =
+
-        new PeerDiscoveryAgent(
+    final MockPeerDiscoveryAgent peerDiscoveryAgent2 = helper.startDiscoveryAgent(peer);
-            vertx,
+    peerDiscoveryAgent2.start().join();
-            keyPair1,
-            DiscoveryConfiguration.create().setBindHost("127.0.0.1").setBindPort(0),
-            () -> true,
-            new PeerBlacklist(),
-            new NodeWhitelistController(PermissioningConfiguration.createDefault()));
-    peerDiscoveryAgent1.start(0).join();
-    final DefaultPeer peer = peerDiscoveryAgent1.getAdvertisedPeer();
-
-    final SECP256K1.KeyPair keyPair2 = SECP256K1.KeyPair.generate();
-    final PeerDiscoveryAgent peerDiscoveryAgent2 =
-        new PeerDiscoveryAgent(
-            vertx,
-            keyPair2,
-            DiscoveryConfiguration.create()
-                .setBindHost("127.0.0.1")
-                .setBindPort(0)
-                .setBootstrapPeers(Lists.newArrayList(peer)),
-            () -> true,
-            new PeerBlacklist(),
-            new NodeWhitelistController(PermissioningConfiguration.createDefault()));
-    peerDiscoveryAgent2.start(0).join();
 
     assertThat(peerDiscoveryAgent2.getPeers().size()).isEqualTo(1);
 
@@ -169,16 +133,17 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
   }
 
   @Test
-  public void doesNotBlacklistPeerForNormalDisconnect() throws Exception {
+  public void doesNotBlacklistPeerForNormalDisconnect() {
     // Start an agent with no bootstrap peers.
     final PeerBlacklist blacklist = new PeerBlacklist();
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList(), blacklist);
+    final MockPeerDiscoveryAgent agent =
+        helper.startDiscoveryAgent(Collections.emptyList(), blacklist);
     // Setup peer
-    final DiscoveryTestSocket peerSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherNode = helper.startDiscoveryAgent();
-    final PeerConnection wirePeer = createAnonymousPeerConnection(peerSocket.getPeer().getId());
+    final PeerConnection wirePeer = createAnonymousPeerConnection(otherNode.getId());
 
     // Bond to peer
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
     assertThat(agent.getPeers()).hasSize(1);
 
     // Disconnect with innocuous reason
@@ -188,23 +153,30 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
     assertThat(agent.getPeers()).hasSize(0);
 
     // Bond again
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
 
     // Check peer was allowed to connect
     assertThat(agent.getPeers()).hasSize(1);
   }
 
+  protected void bondViaIncomingPing(
+      final MockPeerDiscoveryAgent agent, final MockPeerDiscoveryAgent otherNode) {
+    Packet pingPacket = helper.createPingPacket(otherNode, agent);
+    helper.sendMessageBetweenAgents(otherNode, agent, pingPacket);
+  }
+
   @Test
-  public void blacklistPeerForBadBehavior() throws Exception {
+  public void blacklistPeerForBadBehavior() {
     // Start an agent with no bootstrap peers.
     final PeerBlacklist blacklist = new PeerBlacklist();
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList(), blacklist);
+    final MockPeerDiscoveryAgent agent =
+        helper.startDiscoveryAgent(Collections.emptyList(), blacklist);
     // Setup peer
-    final DiscoveryTestSocket peerSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherNode = helper.startDiscoveryAgent();
-    final PeerConnection wirePeer = createAnonymousPeerConnection(peerSocket.getPeer().getId());
+    final PeerConnection wirePeer = createAnonymousPeerConnection(otherNode.getId());
 
     // Bond to peer
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
     assertThat(agent.getPeers()).hasSize(1);
 
     // Disconnect with problematic reason
@@ -214,7 +186,7 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
     assertThat(agent.getPeers()).hasSize(0);
 
     // Bond again
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
 
     // Check peer was not allowed to connect
     assertThat(agent.getPeers()).hasSize(0);
@@ -224,13 +196,14 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
   public void doesNotBlacklistPeerForOurBadBehavior() throws Exception {
     // Start an agent with no bootstrap peers.
     final PeerBlacklist blacklist = new PeerBlacklist();
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList(), blacklist);
+    final MockPeerDiscoveryAgent agent =
+        helper.startDiscoveryAgent(Collections.emptyList(), blacklist);
     // Setup peer
-    final DiscoveryTestSocket peerSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherNode = helper.startDiscoveryAgent();
-    final PeerConnection wirePeer = createAnonymousPeerConnection(peerSocket.getPeer().getId());
+    final PeerConnection wirePeer = createAnonymousPeerConnection(otherNode.getId());
 
     // Bond to peer
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
     assertThat(agent.getPeers()).hasSize(1);
 
     // Disconnect with problematic reason
@@ -240,7 +213,7 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
     assertThat(agent.getPeers()).hasSize(0);
 
     // Bond again
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
 
     // Check peer was allowed to connect
     assertThat(agent.getPeers()).hasSize(1);
@@ -250,13 +223,14 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
   public void blacklistIncompatiblePeer() throws Exception {
     // Start an agent with no bootstrap peers.
     final PeerBlacklist blacklist = new PeerBlacklist();
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList(), blacklist);
+    final MockPeerDiscoveryAgent agent =
+        helper.startDiscoveryAgent(Collections.emptyList(), blacklist);
     // Setup peer
-    final DiscoveryTestSocket peerSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherNode = helper.startDiscoveryAgent();
-    final PeerConnection wirePeer = createAnonymousPeerConnection(peerSocket.getPeer().getId());
+    final PeerConnection wirePeer = createAnonymousPeerConnection(otherNode.getId());
 
     // Bond to peer
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
     assertThat(agent.getPeers()).hasSize(1);
 
     // Disconnect
@@ -266,7 +240,7 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
     assertThat(agent.getPeers()).hasSize(0);
 
     // Bond again
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
 
     // Check peer was not allowed to connect
     assertThat(agent.getPeers()).hasSize(0);
@@ -276,13 +250,14 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
   public void blacklistIncompatiblePeerWhoIssuesDisconnect() throws Exception {
     // Start an agent with no bootstrap peers.
     final PeerBlacklist blacklist = new PeerBlacklist();
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList(), blacklist);
+    final MockPeerDiscoveryAgent agent =
+        helper.startDiscoveryAgent(Collections.emptyList(), blacklist);
     // Setup peer
-    final DiscoveryTestSocket peerSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherNode = helper.startDiscoveryAgent();
-    final PeerConnection wirePeer = createAnonymousPeerConnection(peerSocket.getPeer().getId());
+    final PeerConnection wirePeer = createAnonymousPeerConnection(otherNode.getId());
 
     // Bond to peer
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
     assertThat(agent.getPeers()).hasSize(1);
 
     // Disconnect
@@ -292,7 +267,7 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
     assertThat(agent.getPeers()).hasSize(0);
 
     // Bond again
-    bondViaIncomingPing(agent, peerSocket);
+    bondViaIncomingPing(agent, otherNode);
 
     // Check peer was not allowed to connect
     assertThat(agent.getPeers()).hasSize(0);
@@ -300,20 +275,16 @@ public class PeerDiscoveryAgentTest extends AbstractPeerDiscoveryTest {
 
   @Test
   public void shouldBeActiveWhenConfigIsTrue() {
-    final DiscoveryConfiguration config = new DiscoveryConfiguration();
+    AgentBuilder agentBuilder = helper.agentBuilder().active(true);
-    config.setActive(true).setBindPort(0);
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(agentBuilder);
-
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(config, new PeerBlacklist());
 
     assertThat(agent.isActive()).isTrue();
   }
 
   @Test
   public void shouldNotBeActiveWhenConfigIsFalse() {
-    final DiscoveryConfiguration config = new DiscoveryConfiguration();
+    AgentBuilder agentBuilder = helper.agentBuilder().active(false);
-    config.setActive(false).setBindPort(0);
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(agentBuilder);
-
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(config, new PeerBlacklist());
 
     assertThat(agent.isActive()).isFalse();
   }

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryBondingTest.java

@@ -16,38 +16,43 @@ import static java.util.Collections.emptyList;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.FindNeighborsPacketData;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent.IncomingPacket;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketType;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PingPacketData;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PongPacketData;
 
 import java.util.Collections;
-import java.util.concurrent.TimeUnit;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
 import org.junit.Test;
 
-public class PeerDiscoveryBondingTest extends AbstractPeerDiscoveryTest {
+public class PeerDiscoveryBondingTest {
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
-  public void pongSentUponPing() throws Exception {
+  public void pongSentUponPing() {
     // Start an agent with no bootstrap peers.
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList());
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(Collections.emptyList());
 
     // Start a test peer and send a PING packet to the agent under test.
-    final DiscoveryTestSocket discoveryTestSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherAgent = helper.startDiscoveryAgent();
-
+    final Packet ping = helper.createPingPacket(otherAgent, agent);
-    final PingPacketData ping =
+    helper.sendMessageBetweenAgents(otherAgent, agent, ping);
-        PingPacketData.create(
+
-            discoveryTestSocket.getPeer().getEndpoint(), agent.getAdvertisedPeer().getEndpoint());
+    final List<IncomingPacket> otherAgentIncomingPongs =
-    final Packet packet = Packet.create(PacketType.PING, ping, discoveryTestSocket.getKeyPair());
+        otherAgent
-    discoveryTestSocket.sendToAgent(agent, packet);
+            .getIncomingPackets()
-
+            .stream()
-    final Packet pongPacket = discoveryTestSocket.getIncomingPackets().poll(10, TimeUnit.SECONDS);
+            .filter(p -> p.packet.getType().equals(PacketType.PONG))
-    assertThat(pongPacket.getType()).isEqualTo(PacketType.PONG);
+            .collect(Collectors.toList());
-    assertThat(pongPacket.getPacketData(PongPacketData.class)).isPresent();
+    assertThat(otherAgentIncomingPongs.size()).isEqualTo(1);
-
+
-    final PongPacketData pong = pongPacket.getPacketData(PongPacketData.class).get();
+    final PongPacketData pong =
-    assertThat(pong.getTo()).isEqualTo(discoveryTestSocket.getPeer().getEndpoint());
+        otherAgentIncomingPongs.get(0).packet.getPacketData(PongPacketData.class).get();
+    assertThat(pong.getTo()).isEqualTo(otherAgent.getAdvertisedPeer().getEndpoint());
 
     // The agent considers the test peer BONDED.
     assertThat(agent.getPeers()).hasSize(1);
@@ -57,38 +62,38 @@ public class PeerDiscoveryBondingTest extends AbstractPeerDiscoveryTest {
   @Test
   public void neighborsPacketNotSentUnlessBonded() throws InterruptedException {
     // Start an agent.
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(emptyList());
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(emptyList());
 
     // Start a test peer that will send a FIND_NEIGHBORS to the agent under test. It should be
     // ignored because
     // we haven't bonded.
-    final DiscoveryTestSocket discoveryTestSocket = startTestSocket();
+    final MockPeerDiscoveryAgent otherNode = helper.startDiscoveryAgent();
-    final FindNeighborsPacketData data =
+    final FindNeighborsPacketData data = FindNeighborsPacketData.create(otherNode.getId());
-        FindNeighborsPacketData.create(discoveryTestSocket.getPeer().getId());
+    Packet packet = Packet.create(PacketType.FIND_NEIGHBORS, data, otherNode.getKeyPair());
-    Packet packet =
+    helper.sendMessageBetweenAgents(otherNode, agent, packet);
-        Packet.create(PacketType.FIND_NEIGHBORS, data, discoveryTestSocket.getKeyPair());
-    discoveryTestSocket.sendToAgent(agent, packet);
 
-    // No responses received in 2 seconds.
+    // No responses received
-    final Packet incoming = discoveryTestSocket.getIncomingPackets().poll(2, TimeUnit.SECONDS);
+    final List<IncomingPacket> incoming = otherNode.getIncomingPackets();
-    assertThat(incoming).isNull();
+    assertThat(incoming.size()).isEqualTo(0);
 
     // Create and dispatch a PING packet.
-    final PingPacketData ping =
+    final Packet ping = helper.createPingPacket(otherNode, agent);
-        PingPacketData.create(
+    helper.sendMessageBetweenAgents(otherNode, agent, ping);
-            discoveryTestSocket.getPeer().getEndpoint(), agent.getAdvertisedPeer().getEndpoint());
-    packet = Packet.create(PacketType.PING, ping, discoveryTestSocket.getKeyPair());
-    discoveryTestSocket.sendToAgent(agent, packet);
 
     // Now we received a PONG.
-    final Packet pongPacket = discoveryTestSocket.getIncomingPackets().poll(2, TimeUnit.SECONDS);
+    final List<IncomingPacket> incomingPongs =
-    assertThat(pongPacket.getType()).isEqualTo(PacketType.PONG);
+        otherNode
-    assertThat(pongPacket.getPacketData(PongPacketData.class)).isPresent();
+            .getIncomingPackets()
-
+            .stream()
-    final PongPacketData pong = pongPacket.getPacketData(PongPacketData.class).get();
+            .filter(p -> p.packet.getType().equals(PacketType.PONG))
-    assertThat(pong.getTo()).isEqualTo(discoveryTestSocket.getPeer().getEndpoint());
+            .collect(Collectors.toList());
+    assertThat(incomingPongs.size()).isEqualTo(1);
+    Optional<PongPacketData> maybePongData =
+        incomingPongs.get(0).packet.getPacketData(PongPacketData.class);
+    assertThat(maybePongData).isPresent();
+    assertThat(maybePongData.get().getTo()).isEqualTo(otherNode.getAdvertisedPeer().getEndpoint());
 
     // No more packets.
-    assertThat(discoveryTestSocket.getIncomingPackets()).hasSize(0);
+    assertThat(otherNode.getIncomingPackets()).hasSize(0);
   }
 }

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryBootstrappingTest.java

@@ -16,8 +16,9 @@ import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.awaitility.Awaitility.await;
 
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent.IncomingPacket;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketType;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PingPacketData;
@@ -25,67 +26,71 @@ import tech.pegasys.pantheon.ethereum.p2p.peers.Peer;
 import tech.pegasys.pantheon.util.bytes.BytesValue;
 
 import java.util.List;
-import java.util.concurrent.TimeUnit;
-import java.util.stream.Stream;
 
 import org.junit.Test;
 
-public class PeerDiscoveryBootstrappingTest extends AbstractPeerDiscoveryTest {
+public class PeerDiscoveryBootstrappingTest {
+
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
-  public void bootstrappingPingsSentSingleBootstrapPeer() throws Exception {
+  public void bootstrappingPingsSentSingleBootstrapPeer() {
     // Start one test peer and use it as a bootstrap peer.
-    final DiscoveryTestSocket discoveryTestSocket = startTestSocket();
+    final MockPeerDiscoveryAgent testAgent = helper.startDiscoveryAgent();
-    final List<DiscoveryPeer> bootstrapPeers = singletonList(discoveryTestSocket.getPeer());
 
     // Start an agent.
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(bootstrapPeers);
+    final PeerDiscoveryAgent agent = helper.startDiscoveryAgent(testAgent.getAdvertisedPeer());
-
-    final Packet packet = discoveryTestSocket.getIncomingPackets().poll(2, TimeUnit.SECONDS);
 
-    assertThat(packet.getType()).isEqualTo(PacketType.PING);
+    final List<IncomingPacket> incomingPackets =
-    assertThat(packet.getNodeId()).isEqualTo(agent.getAdvertisedPeer().getId());
+        testAgent
+            .getIncomingPackets()
+            .stream()
+            .filter(p -> p.packet.getType().equals(PacketType.PING))
+            .collect(toList());
+    assertThat(incomingPackets.size()).isEqualTo(1);
+    Packet pingPacket = incomingPackets.get(0).packet;
+    assertThat(pingPacket.getNodeId()).isEqualTo(agent.getAdvertisedPeer().getId());
 
-    final PingPacketData pingData = packet.getPacketData(PingPacketData.class).get();
+    final PingPacketData pingData = pingPacket.getPacketData(PingPacketData.class).get();
     assertThat(pingData.getExpiration())
         .isGreaterThanOrEqualTo(System.currentTimeMillis() / 1000 - 10000);
     assertThat(pingData.getFrom()).isEqualTo(agent.getAdvertisedPeer().getEndpoint());
-    assertThat(pingData.getTo()).isEqualTo(discoveryTestSocket.getPeer().getEndpoint());
+    assertThat(pingData.getTo()).isEqualTo(testAgent.getAdvertisedPeer().getEndpoint());
   }
 
   @Test
   public void bootstrappingPingsSentMultipleBootstrapPeers() {
-    // Start three test peers.
-    startTestSockets(3);
-
     // Use these peers as bootstrap peers.
+    final List<MockPeerDiscoveryAgent> bootstrapAgents = helper.startDiscoveryAgents(3);
     final List<DiscoveryPeer> bootstrapPeers =
-        discoveryTestSockets.stream().map(DiscoveryTestSocket::getPeer).collect(toList());
+        bootstrapAgents.stream().map(PeerDiscoveryAgent::getAdvertisedPeer).collect(toList());
 
     // Start five agents.
-    startDiscoveryAgents(5, bootstrapPeers);
+    List<MockPeerDiscoveryAgent> agents = helper.startDiscoveryAgents(5, bootstrapPeers);
 
     // Assert that all test peers received a Find Neighbors packet.
-    for (final DiscoveryTestSocket peer : discoveryTestSockets) {
+    for (final MockPeerDiscoveryAgent bootstrapAgent : bootstrapAgents) {
       // Five messages per peer (sent by each of the five agents).
-      final List<Packet> packets = Stream.generate(peer::compulsoryPoll).limit(5).collect(toList());
+      final List<Packet> packets =
-
+          bootstrapAgent.getIncomingPackets().stream().map(p -> p.packet).collect(toList());
-      // No more messages left.
-      assertThat(peer.getIncomingPackets().size()).isEqualTo(0);
 
       // Assert that the node IDs we received belong to the test agents.
-      final List<BytesValue> peerIds = packets.stream().map(Packet::getNodeId).collect(toList());
+      final List<BytesValue> senderIds =
-      final List<BytesValue> nodeIds =
+          packets.stream().map(Packet::getNodeId).distinct().collect(toList());
+      final List<BytesValue> agentIds =
           agents
               .stream()
               .map(PeerDiscoveryAgent::getAdvertisedPeer)
               .map(Peer::getId)
+              .distinct()
               .collect(toList());
 
-      assertThat(peerIds).containsExactlyInAnyOrderElementsOf(nodeIds);
+      assertThat(senderIds).containsExactlyInAnyOrderElementsOf(agentIds);
 
-      // Traverse all received packets.
+      // Traverse all received pings.
-      for (final Packet packet : packets) {
+      List<Packet> pingPackets =
+          packets.stream().filter(p -> p.getType().equals(PacketType.PING)).collect(toList());
+      for (final Packet packet : pingPackets) {
         // Assert that the packet was a Find Neighbors one.
         assertThat(packet.getType()).isEqualTo(PacketType.PING);
 
@@ -93,7 +98,7 @@ public class PeerDiscoveryBootstrappingTest extends AbstractPeerDiscoveryTest {
         final PingPacketData ping = packet.getPacketData(PingPacketData.class).get();
         assertThat(ping.getExpiration())
             .isGreaterThanOrEqualTo(System.currentTimeMillis() / 1000 - 10000);
-        assertThat(ping.getTo()).isEqualTo(peer.getPeer().getEndpoint());
+        assertThat(ping.getTo()).isEqualTo(bootstrapAgent.getAdvertisedPeer().getEndpoint());
       }
     }
   }
@@ -101,25 +106,18 @@ public class PeerDiscoveryBootstrappingTest extends AbstractPeerDiscoveryTest {
   @Test
   public void bootstrappingPeersListUpdated() {
     // Start an agent.
-    final PeerDiscoveryAgent bootstrapAgent = startDiscoveryAgent(emptyList());
+    final PeerDiscoveryAgent bootstrapAgent = helper.startDiscoveryAgent(emptyList());
 
     // Start other five agents, pointing to the one above as a bootstrap peer.
-    final List<PeerDiscoveryAgent> otherAgents =
+    final List<MockPeerDiscoveryAgent> otherAgents =
-        startDiscoveryAgents(5, singletonList(bootstrapAgent.getAdvertisedPeer()));
+        helper.startDiscoveryAgents(5, singletonList(bootstrapAgent.getAdvertisedPeer()));
 
     final BytesValue[] otherPeersIds =
-        otherAgents
+        otherAgents.stream().map(PeerDiscoveryAgent::getId).toArray(BytesValue[]::new);
-            .stream()
+
-            .map(PeerDiscoveryAgent::getAdvertisedPeer)
-            .map(Peer::getId)
-            .toArray(BytesValue[]::new);
-    await()
-        .atMost(5, TimeUnit.SECONDS)
-        .untilAsserted(
-            () ->
     assertThat(bootstrapAgent.getPeers())
         .extracting(Peer::getId)
-                    .containsExactlyInAnyOrder(otherPeersIds));
+        .containsExactlyInAnyOrder(otherPeersIds);
 
     assertThat(bootstrapAgent.getPeers())
         .allMatch(p -> p.getStatus() == PeerDiscoveryStatus.BONDED);
@@ -128,9 +126,7 @@ public class PeerDiscoveryBootstrappingTest extends AbstractPeerDiscoveryTest {
     // and will
     // bond with them, ultimately adding all 7 nodes in the network to its table.
     final PeerDiscoveryAgent newAgent =
-        startDiscoveryAgent(singletonList(bootstrapAgent.getAdvertisedPeer()));
+        helper.startDiscoveryAgent(bootstrapAgent.getAdvertisedPeer());
-    await()
+    assertThat(newAgent.getPeers()).hasSize(6);
-        .atMost(5, TimeUnit.SECONDS)
-        .untilAsserted(() -> assertThat(newAgent.getPeers()).hasSize(6));
   }
 }

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryObserversTest.java

@@ -12,48 +12,41 @@
  */
 package tech.pegasys.pantheon.ethereum.p2p.discovery;
 
-import static io.vertx.core.Vertx.vertx;
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatCode;
-import static org.awaitility.Awaitility.await;
-import static tech.pegasys.pantheon.ethereum.p2p.NetworkingTestHelper.configWithRandomPorts;
 
-import tech.pegasys.pantheon.crypto.SECP256K1;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryEvent.PeerBondedEvent;
-import tech.pegasys.pantheon.ethereum.p2p.peers.Peer;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent;
-import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
+import tech.pegasys.pantheon.util.bytes.BytesValue;
-import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
-import tech.pegasys.pantheon.ethereum.permissioning.PermissioningConfiguration;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.awaitility.core.ConditionTimeoutException;
 import org.junit.Test;
 
-public class PeerDiscoveryObserversTest extends AbstractPeerDiscoveryTest {
+public class PeerDiscoveryObserversTest {
   private static final Logger LOG = LogManager.getLogger();
   private static final int BROADCAST_TCP_PORT = 26422;
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
   public void addAndRemoveObservers() {
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList());
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(Collections.emptyList());
     assertThat(agent.getObserverCount()).isEqualTo(0);
 
     final long id1 = agent.observePeerBondedEvents((event) -> {});
     final long id2 = agent.observePeerBondedEvents((event) -> {});
     final long id3 = agent.observePeerBondedEvents((event) -> {});
-    final long id4 = agent.observePeerDroppedEvents((event) -> {});
+    final long id4 = agent.observePeerBondedEvents((event) -> {});
-    final long id5 = agent.observePeerDroppedEvents((event) -> {});
+    final long id5 = agent.observePeerBondedEvents((event) -> {});
-    final long id6 = agent.observePeerDroppedEvents((event) -> {});
+    final long id6 = agent.observePeerBondedEvents((event) -> {});
     assertThat(agent.getObserverCount()).isEqualTo(6);
 
     agent.removePeerBondedObserver(id1);
@@ -61,25 +54,25 @@ public class PeerDiscoveryObserversTest extends AbstractPeerDiscoveryTest {
     assertThat(agent.getObserverCount()).isEqualTo(4);
 
     agent.removePeerBondedObserver(id3);
-    agent.removePeerDroppedObserver(id4);
+    agent.removePeerBondedObserver(id4);
     assertThat(agent.getObserverCount()).isEqualTo(2);
 
-    agent.removePeerDroppedObserver(id5);
+    agent.removePeerBondedObserver(id5);
-    agent.removePeerDroppedObserver(id6);
+    agent.removePeerBondedObserver(id6);
     assertThat(agent.getObserverCount()).isEqualTo(0);
 
     final long id7 = agent.observePeerBondedEvents((event) -> {});
-    final long id8 = agent.observePeerDroppedEvents((event) -> {});
+    final long id8 = agent.observePeerBondedEvents((event) -> {});
     assertThat(agent.getObserverCount()).isEqualTo(2);
 
     agent.removePeerBondedObserver(id7);
-    agent.removePeerDroppedObserver(id8);
+    agent.removePeerBondedObserver(id8);
     assertThat(agent.getObserverCount()).isEqualTo(0);
   }
 
   @Test
   public void removeInexistingObserver() {
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList());
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(Collections.emptyList());
     assertThat(agent.getObserverCount()).isEqualTo(0);
 
     agent.observePeerBondedEvents((event) -> {});
@@ -89,14 +82,21 @@ public class PeerDiscoveryObserversTest extends AbstractPeerDiscoveryTest {
   @Test
   public void peerBondedObserverTriggered() throws TimeoutException, InterruptedException {
     // Create 3 discovery agents with no bootstrap peers.
-    final List<PeerDiscoveryAgent> others1 = startDiscoveryAgents(3, Collections.emptyList());
+    final List<MockPeerDiscoveryAgent> others1 =
+        helper.startDiscoveryAgents(3, Collections.emptyList());
     final List<DiscoveryPeer> peers1 =
-        others1.stream().map(PeerDiscoveryAgent::getAdvertisedPeer).collect(Collectors.toList());
+        others1
+            .stream()
+            .map(MockPeerDiscoveryAgent::getAdvertisedPeer)
+            .collect(Collectors.toList());
 
     // Create two discovery agents pointing to the above as bootstrap peers.
-    final List<PeerDiscoveryAgent> others2 = startDiscoveryAgents(2, peers1);
+    final List<MockPeerDiscoveryAgent> others2 = helper.startDiscoveryAgents(2, peers1);
     final List<DiscoveryPeer> peers2 =
-        others2.stream().map(PeerDiscoveryAgent::getAdvertisedPeer).collect(Collectors.toList());
+        others2
+            .stream()
+            .map(MockPeerDiscoveryAgent::getAdvertisedPeer)
+            .collect(Collectors.toList());
 
     // A list of all peers.
     final List<DiscoveryPeer> allPeers = new ArrayList<>(peers1);
@@ -104,42 +104,26 @@ public class PeerDiscoveryObserversTest extends AbstractPeerDiscoveryTest {
 
     // Create a discovery agent (which we'll assert on), using the above two peers as bootstrap
     // peers.
-    final PeerDiscoveryAgent agent =
+    final MockPeerDiscoveryAgent agent = helper.createDiscoveryAgent(peers2);
-        new PeerDiscoveryAgent(
-            vertx(),
-            SECP256K1.KeyPair.generate(),
-            configWithRandomPorts().getDiscovery().setBootstrapPeers(peers2),
-            () -> true,
-            new PeerBlacklist(),
-            new NodeWhitelistController(PermissioningConfiguration.createDefault()));
-
     // A queue for storing peer bonded events.
-    final ArrayBlockingQueue<PeerBondedEvent> queue = new ArrayBlockingQueue<>(10);
+    final List<PeerBondedEvent> events = new ArrayList<>(10);
-    agent.observePeerBondedEvents(queue::add);
+    agent.observePeerBondedEvents(events::add);
-    assertThatCode(() -> agent.start(BROADCAST_TCP_PORT).get(5, TimeUnit.SECONDS))
+    agent.start();
-        .doesNotThrowAnyException();
+
-
+    final HashSet<BytesValue> seenPeers = new HashSet<>();
-    // Wait until we've received 5 events.
+    List<DiscoveryPeer> discoveredPeers =
-    try {
+        events
-      await()
+            .stream()
-          .atMost(5, TimeUnit.SECONDS)
+            .map(PeerDiscoveryEvent::getPeer)
-          .untilAsserted(() -> assertThat(queue.size()).isEqualTo(5));
+            // We emit some duplicate events when the tcp port differs (in terms of presence) for a
-    } catch (final ConditionTimeoutException | AssertionError e) {
+            // peer,
-      final List<String> events = new ArrayList<>();
+            // filter peers by id to remove duplicates (See: DefaultPeer::equals).
-      queue.forEach(evt -> events.add(evt.toString()));
+            // TODO: Should we evaluate peer equality based on id??
-      LOG.error("Queue:\n" + String.join("\n", events), e);
+            .filter((p) -> seenPeers.add(p.getId()))
-      throw e;
+            .collect(Collectors.toList());
-    }
+    assertThat(discoveredPeers.size()).isEqualTo(allPeers.size());
-    // Wait one second and check we've received no more events.
-    Thread.sleep(1000);
-    assertThat(queue.size()).isEqualTo(5);
-
-    // Extract all events and perform asserts on them.
-    final List<PeerBondedEvent> events = new ArrayList<>(5);
-    queue.drainTo(events, 5);
 
-    assertThat(events)
+    assertThat(discoveredPeers)
-        .extracting(PeerDiscoveryEvent::getPeer)
         .extracting(DiscoveryPeer::getId)
         .containsExactlyInAnyOrderElementsOf(
             allPeers.stream().map(DiscoveryPeer::getId).collect(Collectors.toList()));
@@ -149,38 +133,33 @@ public class PeerDiscoveryObserversTest extends AbstractPeerDiscoveryTest {
   @Test
   public void multiplePeerBondedObserversTriggered() {
     // Create 3 discovery agents with no bootstrap peers.
-    final List<PeerDiscoveryAgent> others = startDiscoveryAgents(3, Collections.emptyList());
+    final List<MockPeerDiscoveryAgent> others =
-    final Peer peer = others.stream().map(PeerDiscoveryAgent::getAdvertisedPeer).findFirst().get();
+        helper.startDiscoveryAgents(3, Collections.emptyList());
+    final DiscoveryPeer peer = others.get(0).getAdvertisedPeer();
 
     // Create a discovery agent (which we'll assert on), using the above two peers as bootstrap
     // peers.
-    final PeerDiscoveryAgent agent =
+    final MockPeerDiscoveryAgent agent = helper.createDiscoveryAgent(peer);
-        new PeerDiscoveryAgent(
-            vertx(),
-            SECP256K1.KeyPair.generate(),
-            configWithRandomPorts()
-                .getDiscovery()
-                .setBootstrapPeers(Collections.singletonList(peer)),
-            () -> true,
-            new PeerBlacklist(),
-            new NodeWhitelistController(PermissioningConfiguration.createDefault()));
 
     // Create 5 queues and subscribe them to peer bonded events.
-    final List<ArrayBlockingQueue<PeerBondedEvent>> queues =
+    final List<List<PeerBondedEvent>> queues =
-        Stream.generate(() -> new ArrayBlockingQueue<PeerBondedEvent>(10))
+        Stream.generate(() -> new ArrayList<PeerBondedEvent>(10))
             .limit(5)
             .collect(Collectors.toList());
     queues.forEach(q -> agent.observePeerBondedEvents(q::add));
 
     // Start the agent and wait until each queue receives one event.
-    agent.start(BROADCAST_TCP_PORT);
+    agent.start();
-    await()
+    for (List<PeerBondedEvent> eventQueue : queues) {
-        .atMost(5, TimeUnit.SECONDS)
+      assertThat(eventQueue.size()).isEqualTo(1);
-        .untilAsserted(() -> assertThat(queues).allMatch(q -> q.size() == 1));
+    }
 
     // All events are for the same peer.
     final List<PeerBondedEvent> events =
-        queues.stream().map(ArrayBlockingQueue::poll).collect(Collectors.toList());
+        Stream.of(queues)
+            .flatMap(Collection::stream)
+            .flatMap(Collection::stream)
+            .collect(Collectors.toList());
     assertThat(events).extracting(PeerDiscoveryEvent::getPeer).allMatch(p -> p.equals(peer));
 
     // We can event check that the event instance is the same across all queues.

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryPacketSedesTest.java

@@ -15,8 +15,6 @@ package tech.pegasys.pantheon.ethereum.p2p.discovery;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.data.Offset.offset;
 import static org.junit.Assert.assertNotNull;
-import static tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper.generateKeyPairs;
-import static tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper.generatePeers;
 
 import tech.pegasys.pantheon.crypto.SECP256K1;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.FindNeighborsPacketData;
@@ -29,7 +27,6 @@ import tech.pegasys.pantheon.ethereum.rlp.RLPException;
 import tech.pegasys.pantheon.util.bytes.BytesValue;
 import tech.pegasys.pantheon.util.bytes.MutableBytesValue;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.Random;
 
@@ -37,6 +34,7 @@ import io.vertx.core.buffer.Buffer;
 import org.junit.Test;
 
 public class PeerDiscoveryPacketSedesTest {
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
   public void serializeDeserializeEntirePacket() {
@@ -79,7 +77,7 @@ public class PeerDiscoveryPacketSedesTest {
 
   @Test
   public void neighborsPacketData() {
-    final List<DiscoveryPeer> peers = Arrays.asList(generatePeers(generateKeyPairs(5)));
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(5);
 
     final NeighborsPacketData packet = NeighborsPacketData.create(peers);
     final BytesValue serialized = RLP.encode(packet::writeTo);

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryTestHelper.java

@@ -13,27 +13,206 @@
 package tech.pegasys.pantheon.ethereum.p2p.discovery;
 
 import tech.pegasys.pantheon.crypto.SECP256K1;
-import tech.pegasys.pantheon.ethereum.p2p.peers.Endpoint;
+import tech.pegasys.pantheon.crypto.SECP256K1.KeyPair;
+import tech.pegasys.pantheon.ethereum.p2p.config.DiscoveryConfiguration;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketType;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PingPacketData;
+import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
+import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
+import tech.pegasys.pantheon.ethereum.permissioning.PermissioningConfiguration;
+import tech.pegasys.pantheon.util.bytes.BytesValue;
 
-import java.util.OptionalInt;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 public class PeerDiscoveryTestHelper {
+  private static final String LOOPBACK_IP_ADDR = "127.0.0.1";
 
-  public static SECP256K1.KeyPair[] generateKeyPairs(final int count) {
+  private final AtomicInteger nextAvailablePort = new AtomicInteger(1);
-    return Stream.generate(SECP256K1.KeyPair::generate)
+  Map<BytesValue, MockPeerDiscoveryAgent> agents = new HashMap<>();
+
+  public static List<SECP256K1.KeyPair> generateKeyPairs(final int count) {
+    return Stream.generate(SECP256K1.KeyPair::generate).limit(count).collect(Collectors.toList());
+  }
+
+  /**
+   * Starts multiple discovery agents from generated peers.
+   *
+   * @param count the number of agents to start
+   * @return a list of discovery agents.
+   */
+  public List<DiscoveryPeer> createDiscoveryPeers(final int count) {
+    return Stream.generate(this::createDiscoveryPeer).limit(count).collect(Collectors.toList());
+  }
+
+  public List<DiscoveryPeer> createDiscoveryPeers(final List<KeyPair> keyPairs) {
+    return keyPairs.stream().map(this::createDiscoveryPeer).collect(Collectors.toList());
+  }
+
+  public DiscoveryPeer createDiscoveryPeer() {
+    return createDiscoveryPeer(KeyPair.generate());
+  }
+
+  public DiscoveryPeer createDiscoveryPeer(final KeyPair keyPair) {
+    final BytesValue peerId = keyPair.getPublicKey().getEncodedBytes();
+    final int port = nextAvailablePort.incrementAndGet();
+    return new DiscoveryPeer(peerId, LOOPBACK_IP_ADDR, port, port);
+  }
+
+  public Packet createPingPacket(
+      final MockPeerDiscoveryAgent fromAgent, final MockPeerDiscoveryAgent toAgent) {
+    return Packet.create(
+        PacketType.PING,
+        PingPacketData.create(
+            fromAgent.getAdvertisedPeer().getEndpoint(), toAgent.getAdvertisedPeer().getEndpoint()),
+        fromAgent.getKeyPair());
+  }
+
+  public AgentBuilder agentBuilder() {
+    return new AgentBuilder(agents, nextAvailablePort);
+  }
+
+  public void sendMessageBetweenAgents(
+      final MockPeerDiscoveryAgent fromAgent,
+      final MockPeerDiscoveryAgent toAgent,
+      final Packet packet) {
+    toAgent.processIncomingPacket(fromAgent, packet);
+  }
+
+  /**
+   * Starts multiple discovery agents with the provided boostrap peers.
+   *
+   * @param count the number of agents to start
+   * @param bootstrapPeers the list of bootstrap peers
+   * @return a list of discovery agents.
+   */
+  public List<MockPeerDiscoveryAgent> startDiscoveryAgents(
+      final int count, final List<DiscoveryPeer> bootstrapPeers) {
+    return Stream.generate(() -> startDiscoveryAgent(bootstrapPeers))
         .limit(count)
-        .toArray(SECP256K1.KeyPair[]::new);
+        .collect(Collectors.toList());
   }
 
-  public static DiscoveryPeer[] generatePeers(final SECP256K1.KeyPair... keypairs) {
+  public List<MockPeerDiscoveryAgent> startDiscoveryAgents(final int count) {
-    return Stream.of(keypairs)
+    return Stream.generate(() -> startDiscoveryAgent(Collections.emptyList()))
-        .map(kp -> kp.getPublicKey().getEncodedBytes())
+        .limit(count)
-        .map(bytes -> new DiscoveryPeer(bytes, new Endpoint("127.0.0.1", 1, OptionalInt.empty())))
+        .collect(Collectors.toList());
-        .toArray(DiscoveryPeer[]::new);
   }
 
-  public static DiscoveryPeer[] generateDiscoveryPeers(final SECP256K1.KeyPair... keypairs) {
+  /**
-    return Stream.of(generatePeers(keypairs)).map(DiscoveryPeer::new).toArray(DiscoveryPeer[]::new);
+   * Start a single discovery agent with the provided bootstrap peers.
+   *
+   * @param bootstrapPeers the list of bootstrap peers
+   * @return a list of discovery agents.
+   */
+  public MockPeerDiscoveryAgent startDiscoveryAgent(final List<DiscoveryPeer> bootstrapPeers) {
+    AgentBuilder agentBuilder = agentBuilder().bootstrapPeers(bootstrapPeers);
+
+    return startDiscoveryAgent(agentBuilder);
+  }
+
+  public MockPeerDiscoveryAgent startDiscoveryAgent(final DiscoveryPeer... bootstrapPeers) {
+    AgentBuilder agentBuilder = agentBuilder().bootstrapPeers(bootstrapPeers);
+
+    return startDiscoveryAgent(agentBuilder);
+  }
+
+  /**
+   * Start a single discovery agent with the provided bootstrap peers.
+   *
+   * @param bootstrapPeers the list of bootstrap peers
+   * @param blacklist the peer blacklist
+   * @return a list of discovery agents.
+   */
+  public MockPeerDiscoveryAgent startDiscoveryAgent(
+      final List<DiscoveryPeer> bootstrapPeers, final PeerBlacklist blacklist) {
+    AgentBuilder agentBuilder = agentBuilder().bootstrapPeers(bootstrapPeers).blacklist(blacklist);
+
+    return startDiscoveryAgent(agentBuilder);
+  }
+
+  public MockPeerDiscoveryAgent startDiscoveryAgent(final AgentBuilder agentBuilder) {
+    final MockPeerDiscoveryAgent agent = createDiscoveryAgent(agentBuilder);
+    agent.start();
+    return agent;
+  }
+
+  public MockPeerDiscoveryAgent createDiscoveryAgent(final List<DiscoveryPeer> bootstrapPeers) {
+    AgentBuilder agentBuilder = agentBuilder().bootstrapPeers(bootstrapPeers);
+
+    return createDiscoveryAgent(agentBuilder);
+  }
+
+  public MockPeerDiscoveryAgent createDiscoveryAgent(final DiscoveryPeer... bootstrapPeers) {
+    AgentBuilder agentBuilder = agentBuilder().bootstrapPeers(bootstrapPeers);
+
+    return createDiscoveryAgent(agentBuilder);
+  }
+
+  public MockPeerDiscoveryAgent createDiscoveryAgent(final AgentBuilder agentBuilder) {
+    final MockPeerDiscoveryAgent agent = agentBuilder.build();
+    agents.put(agent.getId(), agent);
+    return agent;
+  }
+
+  public static class AgentBuilder {
+    private final Map<BytesValue, MockPeerDiscoveryAgent> agents;
+    private final AtomicInteger nextAvailablePort;
+
+    private PeerBlacklist blacklist = new PeerBlacklist();
+    private NodeWhitelistController whitelist =
+        new NodeWhitelistController(PermissioningConfiguration.createDefault());
+    private List<DiscoveryPeer> bootstrapPeers = Collections.emptyList();
+    private boolean active = true;
+
+    public AgentBuilder(
+        final Map<BytesValue, MockPeerDiscoveryAgent> agents,
+        final AtomicInteger nextAvailablePort) {
+      this.agents = agents;
+      this.nextAvailablePort = nextAvailablePort;
+    }
+
+    public AgentBuilder bootstrapPeers(final List<DiscoveryPeer> peers) {
+      this.bootstrapPeers = peers;
+      return this;
+    }
+
+    public AgentBuilder bootstrapPeers(final DiscoveryPeer... peers) {
+      this.bootstrapPeers = Arrays.asList(peers);
+      return this;
+    }
+
+    public AgentBuilder whiteList(final NodeWhitelistController whitelist) {
+      this.whitelist = whitelist;
+      return this;
+    }
+
+    public AgentBuilder blacklist(final PeerBlacklist blacklist) {
+      this.blacklist = blacklist;
+      return this;
+    }
+
+    public AgentBuilder active(final boolean active) {
+      this.active = active;
+      return this;
+    }
+
+    public MockPeerDiscoveryAgent build() {
+      final DiscoveryConfiguration config = new DiscoveryConfiguration();
+      config.setBootstrapPeers(bootstrapPeers);
+      config.setBindPort(nextAvailablePort.incrementAndGet());
+      config.setActive(active);
+
+      return new MockPeerDiscoveryAgent(
+          SECP256K1.KeyPair.generate(), config, () -> true, blacklist, whitelist, agents);
+    }
   }
 }

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/PeerDiscoveryTimestampsTest.java

@@ -13,11 +13,13 @@
 package tech.pegasys.pantheon.ethereum.p2p.discovery;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.awaitility.Awaitility.await;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-import tech.pegasys.pantheon.crypto.SECP256K1;
+import tech.pegasys.pantheon.crypto.SECP256K1.KeyPair;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockPeerDiscoveryAgent;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.MockTimerUtil;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.OutboundMessageHandler;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.Packet;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PacketType;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerDiscoveryController;
@@ -26,101 +28,88 @@ import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PingPacketData;
 import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
 import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
 import tech.pegasys.pantheon.ethereum.permissioning.PermissioningConfiguration;
+import tech.pegasys.pantheon.util.Subscribers;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
-import io.vertx.core.Vertx;
 import org.junit.Test;
 
-public class PeerDiscoveryTimestampsTest extends AbstractPeerDiscoveryTest {
+public class PeerDiscoveryTimestampsTest {
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
   public void lastSeenAndFirstDiscoveredTimestampsUpdatedOnMessage() {
     // peer[0] => controller // peer[1] => sender
-    final SECP256K1.KeyPair[] keypairs = PeerDiscoveryTestHelper.generateKeyPairs(2);
+    final List<KeyPair> keypairs = PeerDiscoveryTestHelper.generateKeyPairs(2);
-    final DiscoveryPeer[] peers = PeerDiscoveryTestHelper.generateDiscoveryPeers(keypairs);
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(keypairs);
 
-    final PeerDiscoveryAgent agent = mock(PeerDiscoveryAgent.class);
+    final MockPeerDiscoveryAgent agent = mock(MockPeerDiscoveryAgent.class);
-    when(agent.getAdvertisedPeer()).thenReturn(peers[0]);
+    when(agent.getAdvertisedPeer()).thenReturn(peers.get(0));
+    DiscoveryPeer localPeer = peers.get(0);
+    KeyPair localKeyPair = keypairs.get(0);
 
     final PeerDiscoveryController controller =
         new PeerDiscoveryController(
-            mock(Vertx.class),
+            localKeyPair,
-            agent,
+            localPeer,
             new PeerTable(agent.getAdvertisedPeer().getId()),
             Collections.emptyList(),
+            OutboundMessageHandler.NOOP,
+            new MockTimerUtil(),
             TimeUnit.HOURS.toMillis(1),
             () -> true,
             new PeerBlacklist(),
-            new NodeWhitelistController(PermissioningConfiguration.createDefault()));
+            new NodeWhitelistController(PermissioningConfiguration.createDefault()),
+            new Subscribers<>());
     controller.start();
 
     final PingPacketData ping =
-        PingPacketData.create(peers[1].getEndpoint(), peers[0].getEndpoint());
+        PingPacketData.create(peers.get(1).getEndpoint(), peers.get(0).getEndpoint());
-    final Packet packet = Packet.create(PacketType.PING, ping, keypairs[1]);
+    final Packet packet = Packet.create(PacketType.PING, ping, keypairs.get(1));
 
-    controller.onMessage(packet, peers[1]);
+    controller.onMessage(packet, peers.get(1));
 
     final AtomicLong lastSeen = new AtomicLong();
     final AtomicLong firstDiscovered = new AtomicLong();
 
-    await()
-        .atMost(1, TimeUnit.SECONDS)
-        .untilAsserted(
-            () -> {
     assertThat(controller.getPeers()).hasSize(1);
 
-              final DiscoveryPeer p = controller.getPeers().iterator().next();
+    DiscoveryPeer p = controller.getPeers().iterator().next();
     assertThat(p.getLastSeen()).isGreaterThan(0);
     assertThat(p.getFirstDiscovered()).isGreaterThan(0);
 
     lastSeen.set(p.getLastSeen());
     firstDiscovered.set(p.getFirstDiscovered());
-            });
 
-    controller.onMessage(packet, peers[1]);
+    controller.onMessage(packet, peers.get(1));
 
-    await()
-        .atMost(1, TimeUnit.SECONDS)
-        .untilAsserted(
-            () -> {
     assertThat(controller.getPeers()).hasSize(1);
 
-              final DiscoveryPeer p = controller.getPeers().iterator().next();
+    p = controller.getPeers().iterator().next();
     assertThat(p.getLastSeen()).isGreaterThan(lastSeen.get());
     assertThat(p.getFirstDiscovered()).isEqualTo(firstDiscovered.get());
-            });
   }
 
   @Test
   public void lastContactedTimestampUpdatedOnOutboundMessage() {
-    final PeerDiscoveryAgent agent = startDiscoveryAgent(Collections.emptyList());
+    final MockPeerDiscoveryAgent agent = helper.startDiscoveryAgent(Collections.emptyList());
     assertThat(agent.getPeers()).hasSize(0);
 
     // Start a test peer and send a PING packet to the agent under test.
-    final DiscoveryTestSocket discoveryTestSocket = startTestSocket();
+    final MockPeerDiscoveryAgent testAgent = helper.startDiscoveryAgent();
+    final Packet ping = helper.createPingPacket(testAgent, agent);
+    helper.sendMessageBetweenAgents(testAgent, agent, ping);
 
-    final PingPacketData ping =
+    assertThat(agent.getPeers()).hasSize(1);
-        PingPacketData.create(
-            discoveryTestSocket.getPeer().getEndpoint(), agent.getAdvertisedPeer().getEndpoint());
-    final Packet packet = Packet.create(PacketType.PING, ping, discoveryTestSocket.getKeyPair());
-    discoveryTestSocket.sendToAgent(agent, packet);
-
-    await()
-        .atMost(1, TimeUnit.SECONDS)
-        .untilAsserted(() -> assertThat(agent.getPeers()).hasSize(1));
 
     final AtomicLong lastContacted = new AtomicLong();
     final AtomicLong lastSeen = new AtomicLong();
     final AtomicLong firstDiscovered = new AtomicLong();
 
-    await()
+    DiscoveryPeer peer = agent.getPeers().iterator().next();
-        .atMost(1, TimeUnit.SECONDS)
-        .untilAsserted(
-            () -> {
-              final DiscoveryPeer peer = agent.getPeers().iterator().next();
     final long lc = peer.getLastContacted();
     final long ls = peer.getLastSeen();
     final long fd = peer.getFirstDiscovered();
@@ -132,20 +121,14 @@ public class PeerDiscoveryTimestampsTest extends AbstractPeerDiscoveryTest {
     lastContacted.set(lc);
     lastSeen.set(ls);
     firstDiscovered.set(fd);
-            });
 
     // Send another packet and ensure that timestamps are updated accordingly.
-    discoveryTestSocket.sendToAgent(agent, packet);
+    helper.sendMessageBetweenAgents(testAgent, agent, ping);
 
-    await()
+    peer = agent.getPeers().iterator().next();
-        .atMost(1, TimeUnit.SECONDS)
-        .untilAsserted(
-            () -> {
-              final DiscoveryPeer peer = agent.getPeers().iterator().next();
 
     assertThat(peer.getLastContacted()).isGreaterThan(lastContacted.get());
     assertThat(peer.getLastSeen()).isGreaterThan(lastSeen.get());
     assertThat(peer.getFirstDiscovered()).isEqualTo(firstDiscovered.get());
-            });
   }
 }

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/BucketTest.java

@@ -15,26 +15,26 @@ package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
 import static junit.framework.TestCase.assertFalse;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
-import static tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper.generateDiscoveryPeers;
-import static tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper.generateKeyPairs;
 
 import tech.pegasys.pantheon.ethereum.p2p.discovery.DiscoveryPeer;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper;
 
+import java.util.List;
 import java.util.Optional;
-import java.util.stream.Stream;
 
 import org.junit.Test;
 
 public class BucketTest {
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
   public void successfulAddAndGet() {
     final Bucket kBucket = new Bucket(16);
-    final DiscoveryPeer[] peers = generateDiscoveryPeers(generateKeyPairs(10));
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(10);
-    for (int i = 0; i < peers.length - 1; i++) {
+    for (int i = 0; i < peers.size() - 1; i++) {
-      kBucket.add(peers[i]);
+      kBucket.add(peers.get(i));
     }
-    final DiscoveryPeer testPeer = peers[peers.length - 1];
+    final DiscoveryPeer testPeer = peers.get(peers.size() - 1);
     kBucket.add(testPeer);
     assertThat(testPeer).isEqualTo(kBucket.getAndTouch(testPeer.getId()).get());
   }
@@ -42,48 +42,48 @@ public class BucketTest {
   @Test
   public void unsuccessfulAdd() {
     final Bucket kBucket = new Bucket(16);
-    final DiscoveryPeer[] peers = generateDiscoveryPeers(generateKeyPairs(17));
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(17);
-    for (int i = 0; i < peers.length - 1; i++) {
+    for (int i = 0; i < peers.size() - 1; i++) {
-      kBucket.add(peers[i]);
+      kBucket.add(peers.get(i));
     }
-    final DiscoveryPeer testPeer = peers[peers.length - 1];
+    final DiscoveryPeer testPeer = peers.get(peers.size() - 1);
     final Optional<DiscoveryPeer> evictionCandidate = kBucket.add(testPeer);
-    assertThat(evictionCandidate.get()).isEqualTo(kBucket.getAndTouch(peers[0].getId()).get());
+    assertThat(evictionCandidate.get()).isEqualTo(kBucket.getAndTouch(peers.get(0).getId()).get());
   }
 
   @Test
   public void movedToHead() {
     final Bucket kBucket = new Bucket(16);
-    final DiscoveryPeer[] peers = generateDiscoveryPeers(generateKeyPairs(5));
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(5);
     for (final DiscoveryPeer peer : peers) {
       kBucket.add(peer);
     }
-    kBucket.getAndTouch(peers[0].getId());
+    kBucket.getAndTouch(peers.get(0).getId());
-    assertThat(kBucket.peers().indexOf(peers[0])).isEqualTo(0);
+    assertThat(kBucket.peers().indexOf(peers.get(0))).isEqualTo(0);
   }
 
   @Test
   public void evictPeer() {
     final Bucket kBucket = new Bucket(16);
-    final DiscoveryPeer[] peers = generateDiscoveryPeers(generateKeyPairs(5));
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(5);
     for (final DiscoveryPeer p : peers) {
       kBucket.add(p);
     }
-    kBucket.evict(peers[4]);
+    kBucket.evict(peers.get(4));
-    assertFalse(kBucket.peers().contains(peers[4]));
+    assertFalse(kBucket.peers().contains(peers.get(4)));
   }
 
   @Test
   public void allActionsOnBucket() {
     final Bucket kBucket = new Bucket(16);
-    final DiscoveryPeer[] peers = generateDiscoveryPeers(generateKeyPairs(30));
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(30);
 
     // Try to evict a peer on an empty bucket.
-    assertThat(kBucket.evict(peers[29])).isFalse();
+    assertThat(kBucket.evict(peers.get(29))).isFalse();
 
     // Add the first 16 peers to the bucket.
-    Stream.of(peers)
+    peers
-        .limit(16)
+        .subList(0, 16)
         .forEach(
             p -> {
               assertThat(kBucket.getAndTouch(p.getId())).isNotPresent();
@@ -93,42 +93,57 @@ public class BucketTest {
             });
 
     // Ensure the peer is not there already.
-    assertThat(kBucket.getAndTouch(peers[16].getId())).isNotPresent();
+    assertThat(kBucket.getAndTouch(peers.get(16).getId())).isNotPresent();
 
     // Try to add a 17th peer and check that the eviction candidate matches the first peer.
-    final Optional<DiscoveryPeer> evictionCandidate = kBucket.add(peers[16]);
+    final Optional<DiscoveryPeer> evictionCandidate = kBucket.add(peers.get(16));
-    assertThat(evictionCandidate).isPresent().get().isEqualTo(peers[0]);
+    assertThat(evictionCandidate).isPresent().get().isEqualTo(peers.get(0));
 
     // Try to add a peer that already exists, and check that the bucket size still remains capped at
     // 16.
-    assertThatThrownBy(() -> kBucket.add(peers[0])).isInstanceOf(IllegalArgumentException.class);
+    assertThatThrownBy(() -> kBucket.add(peers.get(0)))
+        .isInstanceOf(IllegalArgumentException.class);
     assertThat(kBucket.peers()).hasSize(16);
 
     // Try to evict a peer that doesn't exist, and check the result is false.
-    assertThat(kBucket.evict(peers[17])).isFalse();
+    assertThat(kBucket.evict(peers.get(17))).isFalse();
     assertThat(kBucket.peers()).hasSize(16);
 
     // Evict a peer from head, another from the middle, and the tail.
-    assertThat(kBucket.evict(peers[0])).isTrue();
+    assertThat(kBucket.evict(peers.get(0))).isTrue();
     assertThat(kBucket.peers()).hasSize(15);
-    assertThat(kBucket.evict(peers[7])).isTrue();
+    assertThat(kBucket.evict(peers.get(7))).isTrue();
     assertThat(kBucket.peers()).hasSize(14);
-    assertThat(kBucket.evict(peers[15])).isTrue();
+    assertThat(kBucket.evict(peers.get(15))).isTrue();
     assertThat(kBucket.peers()).hasSize(13);
 
     // Check that we can now add peers again.
-    assertThat(kBucket.add(peers[0])).isNotPresent();
+    assertThat(kBucket.add(peers.get(0))).isNotPresent();
-    assertThat(kBucket.add(peers[7])).isNotPresent();
+    assertThat(kBucket.add(peers.get(7))).isNotPresent();
-    assertThat(kBucket.add(peers[15])).isNotPresent();
+    assertThat(kBucket.add(peers.get(15))).isNotPresent();
-    assertThat(kBucket.add(peers[17])).isPresent().get().isEqualTo(peers[1]);
+    assertThat(kBucket.add(peers.get(17))).isPresent().get().isEqualTo(peers.get(1));
 
     // Test the touch behaviour.
-    assertThat(kBucket.getAndTouch(peers[6].getId())).isPresent().get().isEqualTo(peers[6]);
+    assertThat(kBucket.getAndTouch(peers.get(6).getId())).isPresent().get().isEqualTo(peers.get(6));
-    assertThat(kBucket.getAndTouch(peers[9].getId())).isPresent().get().isEqualTo(peers[9]);
+    assertThat(kBucket.getAndTouch(peers.get(9).getId())).isPresent().get().isEqualTo(peers.get(9));
 
     assertThat(kBucket.peers())
         .containsSequence(
-            peers[9], peers[6], peers[15], peers[7], peers[0], peers[14], peers[13], peers[12],
+            peers.get(9),
-            peers[11], peers[10], peers[8], peers[5], peers[4], peers[3], peers[2], peers[1]);
+            peers.get(6),
+            peers.get(15),
+            peers.get(7),
+            peers.get(0),
+            peers.get(14),
+            peers.get(13),
+            peers.get(12),
+            peers.get(11),
+            peers.get(10),
+            peers.get(8),
+            peers.get(5),
+            peers.get(4),
+            peers.get(3),
+            peers.get(2),
+            peers.get(1));
   }
 }

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/MockPeerDiscoveryAgent.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
+
+import tech.pegasys.pantheon.crypto.SECP256K1.KeyPair;
+import tech.pegasys.pantheon.ethereum.p2p.config.DiscoveryConfiguration;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.DiscoveryPeer;
+import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryAgent;
+import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
+import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
+import tech.pegasys.pantheon.util.bytes.BytesValue;
+
+import java.net.InetSocketAddress;
+import java.util.ArrayDeque;
+import java.util.Arrays;
+import java.util.Deque;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+
+public class MockPeerDiscoveryAgent extends PeerDiscoveryAgent {
+  // The set of known agents operating on the network
+  private final Map<BytesValue, MockPeerDiscoveryAgent> agentNetwork;
+  private final Deque<IncomingPacket> incomingPackets = new ArrayDeque<>();
+
+  public MockPeerDiscoveryAgent(
+      final KeyPair keyPair,
+      final DiscoveryConfiguration config,
+      final PeerRequirement peerRequirement,
+      final PeerBlacklist peerBlacklist,
+      final NodeWhitelistController nodeWhitelistController,
+      final Map<BytesValue, MockPeerDiscoveryAgent> agentNetwork) {
+    super(keyPair, config, peerRequirement, peerBlacklist, nodeWhitelistController);
+    this.agentNetwork = agentNetwork;
+  }
+
+  public void processIncomingPacket(final MockPeerDiscoveryAgent fromAgent, final Packet packet) {
+    // Cycle packet through encode / decode to make clone of any data
+    // This ensures that any data passed between agents is not shared
+    final Packet packetClone = Packet.decode(packet.encode());
+    incomingPackets.add(new IncomingPacket(fromAgent, packetClone));
+    handleIncomingPacket(fromAgent.getAdvertisedPeer().getEndpoint(), packetClone);
+  }
+
+  /**
+   * Get and clear the list of any incoming packets to this agent.
+   *
+   * @return A list of packets received by this agent
+   */
+  public List<IncomingPacket> getIncomingPackets() {
+    List<IncomingPacket> packets = Arrays.asList(incomingPackets.toArray(new IncomingPacket[0]));
+    incomingPackets.clear();
+    return packets;
+  }
+
+  @Override
+  protected CompletableFuture<InetSocketAddress> listenForConnections() {
+    // Skip network setup for tests
+    InetSocketAddress address =
+        new InetSocketAddress(config.getAdvertisedHost(), config.getBindPort());
+    return CompletableFuture.completedFuture(address);
+  }
+
+  @Override
+  protected CompletableFuture<Void> sendOutgoingPacket(
+      final DiscoveryPeer toPeer, final Packet packet) {
+    CompletableFuture<Void> result = new CompletableFuture<>();
+    MockPeerDiscoveryAgent toAgent = agentNetwork.get(toPeer.getId());
+    if (toAgent == null) {
+      result.completeExceptionally(
+          new Exception(
+              "Attempt to send to unknown peer.  Agents must be constructed through PeerDiscoveryTestHelper."));
+    } else {
+      toAgent.processIncomingPacket(this, packet);
+      result.complete(null);
+    }
+    return result;
+  }
+
+  @Override
+  protected TimerUtil createTimer() {
+    return new MockTimerUtil();
+  }
+
+  @Override
+  public CompletableFuture<?> stop() {
+    return CompletableFuture.completedFuture(null);
+  }
+
+  public KeyPair getKeyPair() {
+    return keyPair;
+  }
+
+  public static class IncomingPacket {
+    public final MockPeerDiscoveryAgent fromAgent;
+    public final Packet packet;
+
+    public IncomingPacket(final MockPeerDiscoveryAgent fromAgent, final Packet packet) {
+      this.fromAgent = fromAgent;
+      this.packet = packet;
+    }
+  }
+}

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/MockTimerUtil.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
+
+public class MockTimerUtil implements TimerUtil {
+  private final AtomicLong nextId = new AtomicLong(0);
+  private final Map<Long, TimerHandler> timerHandlers = new HashMap<>();
+  private final Map<Long, TimerHandler> periodicHandlers = new HashMap<>();
+
+  @Override
+  public long setPeriodic(final long delay, final TimerHandler handler) {
+    long id = nextId.incrementAndGet();
+    periodicHandlers.put(id, handler);
+    return id;
+  }
+
+  @Override
+  public long setTimer(final long delay, final TimerHandler handler) {
+    long id = nextId.incrementAndGet();
+    timerHandlers.put(id, handler);
+    return id;
+  }
+
+  @Override
+  public void cancelTimer(final long timerId) {
+    timerHandlers.remove(timerId);
+    periodicHandlers.remove(timerId);
+  }
+
+  public void runHandlers() {
+    runTimerHandlers();
+    runPeriodicHandlers();
+  }
+
+  public void runTimerHandlers() {
+    // Create a copy of the handlers to avoid concurrent modification as handlers run
+    List<TimerHandler> handlers = new ArrayList<>();
+    timerHandlers.forEach((id, handler) -> handlers.add(handler));
+    timerHandlers.clear();
+
+    handlers.forEach(TimerHandler::handle);
+  }
+
+  public void runPeriodicHandlers() {
+    // Create a copy of the handlers to avoid concurrent modification as handlers run
+    List<TimerHandler> handlers = new ArrayList<>();
+    periodicHandlers.forEach((id, handler) -> handlers.add(handler));
+
+    handlers.forEach(TimerHandler::handle);
+  }
+}

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/PeerDiscoveryTableRefreshTest.java

@@ -14,86 +14,98 @@ package tech.pegasys.pantheon.ethereum.p2p.discovery.internal;
 
 import static java.util.Collections.emptyList;
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.awaitility.Awaitility.await;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 import tech.pegasys.pantheon.crypto.SECP256K1;
+import tech.pegasys.pantheon.crypto.SECP256K1.KeyPair;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.DiscoveryPeer;
-import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryAgent;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper;
 import tech.pegasys.pantheon.ethereum.p2p.peers.PeerBlacklist;
 import tech.pegasys.pantheon.ethereum.p2p.permissioning.NodeWhitelistController;
 import tech.pegasys.pantheon.ethereum.permissioning.PermissioningConfiguration;
+import tech.pegasys.pantheon.util.Subscribers;
 import tech.pegasys.pantheon.util.bytes.BytesValue;
 
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
-import java.util.concurrent.TimeUnit;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
-import io.vertx.core.Vertx;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 
 public class PeerDiscoveryTableRefreshTest {
-  private final Vertx vertx = spy(Vertx.vertx());
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
   public void tableRefreshSingleNode() {
-    final SECP256K1.KeyPair[] keypairs = PeerDiscoveryTestHelper.generateKeyPairs(2);
+    final List<SECP256K1.KeyPair> keypairs = PeerDiscoveryTestHelper.generateKeyPairs(2);
-    final DiscoveryPeer[] peers = PeerDiscoveryTestHelper.generateDiscoveryPeers(keypairs);
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(keypairs);
+    DiscoveryPeer localPeer = peers.get(0);
+    KeyPair localKeyPair = keypairs.get(0);
 
-    final PeerDiscoveryAgent agent = mock(PeerDiscoveryAgent.class);
+    // Create and start the PeerDiscoveryController
-    when(agent.getAdvertisedPeer()).thenReturn(peers[0]);
+    final OutboundMessageHandler outboundMessageHandler = mock(OutboundMessageHandler.class);
-
+    final MockTimerUtil timer = new MockTimerUtil();
-    // Create and start the PeerDiscoveryController, setting the refresh interval to something
-    // small.
     final PeerDiscoveryController controller =
+        spy(
             new PeerDiscoveryController(
-            vertx,
+                localKeyPair,
-            agent,
+                localPeer,
-            new PeerTable(agent.getAdvertisedPeer().getId()),
+                new PeerTable(localPeer.getId()),
                 emptyList(),
-            100,
+                outboundMessageHandler,
+                timer,
+                0,
                 () -> true,
                 new PeerBlacklist(),
-            new NodeWhitelistController(PermissioningConfiguration.createDefault()));
+                new NodeWhitelistController(PermissioningConfiguration.createDefault()),
+                new Subscribers<>()));
     controller.start();
 
     // Send a PING, so as to add a Peer in the controller.
     final PingPacketData ping =
-        PingPacketData.create(peers[1].getEndpoint(), peers[0].getEndpoint());
+        PingPacketData.create(peers.get(1).getEndpoint(), peers.get(0).getEndpoint());
-    final Packet packet = Packet.create(PacketType.PING, ping, keypairs[1]);
+    final Packet packet = Packet.create(PacketType.PING, ping, keypairs.get(1));
-    controller.onMessage(packet, peers[1]);
+    controller.onMessage(packet, peers.get(1));
 
     // Wait until the controller has added the newly found peer.
-    await()
+    assertThat(controller.getPeers()).hasSize(1);
-        .atMost(1, TimeUnit.SECONDS)
-        .untilAsserted(() -> assertThat(controller.getPeers()).hasSize(1));
 
     // As the controller performs refreshes, it'll send FIND_NEIGHBORS packets with random target
     // IDs every time.
     // We capture the packets so that we can later assert on them.
     // Within 1000ms, there should be ~10 packets. But let's be less ambitious and expect at least
     // 5.
-    final ArgumentCaptor<PacketData> packetDataCaptor = ArgumentCaptor.forClass(PacketData.class);
+    final ArgumentCaptor<Packet> captor = ArgumentCaptor.forClass(Packet.class);
-    verify(agent, timeout(1000).atLeast(5))
+    for (int i = 0; i < 5; i++) {
-        .sendPacket(eq(peers[1]), eq(PacketType.FIND_NEIGHBORS), packetDataCaptor.capture());
+      timer.runPeriodicHandlers();
+    }
+    verify(outboundMessageHandler, atLeast(5)).send(eq(peers.get(1)), captor.capture());
+    List<Packet> capturedFindNeighborsPackets =
+        captor
+            .getAllValues()
+            .stream()
+            .filter(p -> p.getType().equals(PacketType.FIND_NEIGHBORS))
+            .collect(Collectors.toList());
+    assertThat(capturedFindNeighborsPackets.size()).isEqualTo(5);
 
-    // Assert that all packets were FIND_NEIGHBORS packets.
+    // Collect targets from find neighbors packets
     final List<BytesValue> targets = new ArrayList<>();
-    for (final PacketData data : packetDataCaptor.getAllValues()) {
+    for (final Packet captured : capturedFindNeighborsPackets) {
-      assertThat(data).isExactlyInstanceOf(FindNeighborsPacketData.class);
+      Optional<FindNeighborsPacketData> maybeData =
-      final FindNeighborsPacketData fnpd = (FindNeighborsPacketData) data;
+          captured.getPacketData(FindNeighborsPacketData.class);
-      targets.add(fnpd.getTarget());
+      assertThat(maybeData).isPresent();
+      final FindNeighborsPacketData neighborsData = maybeData.get();
+      targets.add(neighborsData.getTarget());
     }
 
-    assertThat(targets.size()).isGreaterThanOrEqualTo(5);
+    assertThat(targets.size()).isEqualTo(5);
 
     // All targets are unique.
     assertThat(targets.size()).isEqualTo(new HashSet<>(targets).size());

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/discovery/internal/PeerTableTest.java

@@ -19,15 +19,17 @@ import tech.pegasys.pantheon.ethereum.p2p.discovery.PeerDiscoveryTestHelper;
 import tech.pegasys.pantheon.ethereum.p2p.discovery.internal.PeerTable.AddResult.Outcome;
 import tech.pegasys.pantheon.ethereum.p2p.peers.Peer;
 
+import java.util.List;
+
 import org.junit.Test;
 
 public class PeerTableTest {
+  private final PeerDiscoveryTestHelper helper = new PeerDiscoveryTestHelper();
 
   @Test
   public void addPeer() {
     final PeerTable table = new PeerTable(Peer.randomId(), 16);
-    final DiscoveryPeer[] peers =
+    final List<DiscoveryPeer> peers = helper.createDiscoveryPeers(5);
-        PeerDiscoveryTestHelper.generateDiscoveryPeers(PeerDiscoveryTestHelper.generateKeyPairs(5));
 
     for (final DiscoveryPeer peer : peers) {
       final PeerTable.AddResult result = table.tryAdd(peer);
@@ -39,9 +41,9 @@ public class PeerTableTest {
 
   @Test
   public void addSelf() {
-    final DiscoveryPeer self = new DiscoveryPeer(Peer.randomId(), "127.0.0.1", 12345, 12345);
+    final DiscoveryPeer localPeer = new DiscoveryPeer(Peer.randomId(), "127.0.0.1", 12345, 12345);
-    final PeerTable table = new PeerTable(self.getId(), 16);
+    final PeerTable table = new PeerTable(localPeer.getId(), 16);
-    final PeerTable.AddResult result = table.tryAdd(self);
+    final PeerTable.AddResult result = table.tryAdd(localPeer);
 
     assertThat(result.getOutcome()).isEqualTo(Outcome.SELF);
     assertThat(table.getAllPeers()).hasSize(0);
@@ -50,9 +52,7 @@ public class PeerTableTest {
   @Test
   public void peerExists() {
     final PeerTable table = new PeerTable(Peer.randomId(), 16);
-    final DiscoveryPeer peer =
+    final DiscoveryPeer peer = helper.createDiscoveryPeer();
-        PeerDiscoveryTestHelper.generateDiscoveryPeers(PeerDiscoveryTestHelper.generateKeyPairs(1))[
-            0];
 
     assertThat(table.tryAdd(peer).getOutcome()).isEqualTo(Outcome.ADDED);
 

pantheon/src/main/java/tech/pegasys/pantheon/Runner.java

@@ -152,6 +152,6 @@ public class Runner implements AutoCloseable {
   }
 
   public int getP2pTcpPort() {
-    return networkRunner.getNetwork().getSelf().getPort();
+    return networkRunner.getNetwork().getLocalPeerInfo().getPort();
   }
 }