Fix tests

Author: xxlabaza

client/src/main/java/io/appulse/epmd/java/client/CommandRegistration.java

@@ -16,7 +16,10 @@
 
 package io.appulse.epmd.java.client;
 
+import static lombok.AccessLevel.PRIVATE;
+
 import java.net.InetAddress;
+import java.util.Map;
 
 import io.appulse.epmd.java.client.exception.EpmdRegistrationException;
 import io.appulse.epmd.java.core.model.request.Registration;
@@ -26,6 +29,7 @@
 import lombok.Builder;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
+import lombok.experimental.FieldDefaults;
 
 /**
  * A command for registration a node in a remote EPMD server.
@@ -34,8 +38,11 @@
  * @author Artem Labazin
  */
 @Slf4j
+@FieldDefaults(level = PRIVATE, makeFinal = true)
 final class CommandRegistration extends CommandAbstract<Registration, RegistrationResult> {
 
+  Map<String, Connection> registered;
+
   /**
    * Constructs the command object.
    *
@@ -46,8 +53,9 @@ final class CommandRegistration extends CommandAbstract<Registration, Registrati
    * @param request the command's request to the remote EPMD server
    */
   @Builder
-  CommandRegistration (InetAddress address, Integer port, Registration request) {
+  CommandRegistration (InetAddress address, Integer port, Registration request, Map<String, Connection> registered) {
     super(address, port, request);
+    this.registered = registered;
   }
 
   @Override
@@ -69,6 +77,7 @@ public RegistrationResult get () {
       throw new EpmdRegistrationException();
     }
 
+    registered.put(request.getName(), connection);
     return result;
   }
 }

client/src/main/java/io/appulse/epmd/java/client/EpmdClient.java

@@ -23,8 +23,8 @@
 import java.io.Closeable;
 import java.net.InetAddress;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
@@ -63,7 +63,7 @@
 @FieldDefaults(level = PRIVATE, makeFinal = true)
 public final class EpmdClient implements Closeable {
 
-  Set<String> registered = ConcurrentHashMap.<String>newKeySet();
+  Map<String, Connection> registered = new ConcurrentHashMap<>();
 
   ExecutorService executor;
 
@@ -137,13 +137,14 @@ public EpmdClient (@NonNull InetAddress address, int port) {
    * @return creation id from EPMD
    */
   public CompletableFuture<RegistrationResult> register (@NonNull Registration request) {
-    if (registered.contains(request.getName())) {
+    if (registered.containsKey(request.getName())) {
       val exception = new EpmdRegistrationNameConflictException(request.getName());
       log.error(exception.getMessage(), exception);
       return FutureUtils.completedExceptionally(exception);
     }
 
     val supplier = CommandRegistration.builder()
+        .registered(registered)
         .address(address)
         .port(port)
         .request(request)
@@ -155,7 +156,6 @@ public CompletableFuture<RegistrationResult> register (@NonNull Registration req
           throw new EpmdRegistrationException(throwable);
         })
         .thenApply(result -> {
-          registered.add(request.getName());
           log.info("'{}' was registered successfully", request.getName());
           return result;
         });
@@ -359,8 +359,11 @@ public CompletableFuture<List<NodeDescription>> getNodes (@NonNull InetAddress n
   @SneakyThrows
   public void close () {
     executor.shutdown();
-    registered.clear();
+
     val terminated = executor.awaitTermination(5, SECONDS);
     log.debug("EPMD was successfully terminated - {}", terminated);
+
+    registered.values().forEach(Connection::close);
+    registered.clear();
   }
 }