improves TransportTest

Signed-off-by: Oleh Dokuka <[email protected]> Signed-off-by: Oleh Dokuka <[email protected]>
rsocket · Nov 12, 2021 · 9460cfc · 9460cfc
1 parent 2989042
commit 9460cfc
Showing 1 changed file with 12 additions and 8 deletions.
diff --git a/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java b/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java
@@ -57,6 +57,7 @@
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
 import reactor.core.Disposable;
+import reactor.core.Disposables;
 import reactor.core.Exceptions;
 import reactor.core.Fuseable;
 import reactor.core.publisher.Flux;
@@ -710,6 +711,7 @@ private static class DisconnectingDuplexConnection implements DuplexConnection {
       private final String tag;
       final DuplexConnection source;
       final Duration delay;
+      final Disposable.Swap disposables = Disposables.swap();
 
       DisconnectingDuplexConnection(String tag, DuplexConnection source, Duration delay) {
         this.tag = tag;
@@ -719,6 +721,7 @@ private static class DisconnectingDuplexConnection implements DuplexConnection {
 
       @Override
       public void dispose() {
+        disposables.dispose();
         source.dispose();
       }
 
@@ -747,14 +750,15 @@ public Flux<ByteBuf> receive() {
                 bb -> {
                   if (!receivedFirst) {
                     receivedFirst = true;
-                    Mono.delay(delay)
-                        .takeUntilOther(source.onClose())
-                        .subscribe(
-                            __ -> {
-                              logger.warn(
-                                  "Tag {}. Disposing Connection[{}]", tag, source.hashCode());
-                              source.dispose();
-                            });
+                    disposables.replace(
+                        Mono.delay(delay)
+                            .takeUntilOther(source.onClose())
+                            .subscribe(
+                                __ -> {
+                                  logger.warn(
+                                      "Tag {}. Disposing Connection[{}]", tag, source.hashCode());
+                                  source.dispose();
+                                }));
                   }
                 });
       }