Make reactive subscription request with Long.MAX_VALUE unbounded and update testkit backend

driver/src/main/java/org/neo4j/driver/internal/cursor/RxResultCursorImpl.java

@@ -88,6 +88,10 @@ public void installRecordConsumer( BiConsumer<Record,Throwable> recordConsumer )
     @Override
     public void request( long n )
     {
+        if ( n == Long.MAX_VALUE )
+        {
+            n = -1;
+        }
         pullHandler.request( n );
     }
 

driver/src/test/java/org/neo4j/driver/internal/cursor/RxResultCursorImplTest.java

@@ -122,6 +122,21 @@ void shouldPull()
         verify( pullHandler ).request( 100 );
     }
 
+    @Test
+    void shouldPullUnboundedOnLongMax()
+    {
+        // Given
+        RunResponseHandler runHandler = newRunResponseHandler();
+        PullResponseHandler pullHandler = mock( PullResponseHandler.class );
+        RxResultCursor cursor = new RxResultCursorImpl( runHandler, pullHandler );
+
+        // When
+        cursor.request( Long.MAX_VALUE );
+
+        // Then
+        verify( pullHandler ).request( -1 );
+    }
+
     @Test
     void shouldCancel()
     {

driver/src/test/java/org/neo4j/driver/internal/reactive/util/ListBasedPullHandler.java

@@ -22,15 +22,15 @@
 import java.util.Map;
 
 import org.neo4j.driver.Query;
+import org.neo4j.driver.Record;
+import org.neo4j.driver.Value;
 import org.neo4j.driver.internal.handlers.PullResponseCompletionListener;
 import org.neo4j.driver.internal.handlers.RunResponseHandler;
 import org.neo4j.driver.internal.handlers.pulln.BasicPullResponseHandler;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.util.MetadataExtractor;
 import org.neo4j.driver.internal.util.QueryKeys;
 import org.neo4j.driver.internal.value.BooleanValue;
-import org.neo4j.driver.Record;
-import org.neo4j.driver.Value;
 import org.neo4j.driver.summary.ResultSummary;
 
 import static java.util.Collections.emptyList;
@@ -81,7 +81,7 @@ private ListBasedPullHandler( List<Record> list, Throwable error )
     public void request( long n )
     {
         super.request( n );
-        while ( index < list.size() && n-- > 0 )
+        while ( index < list.size() && (n == -1 || n-- > 0) )
         {
             onRecord( list.get( index++ ).values().toArray( new Value[0] ) );
         }

testkit-backend/src/main/java/neo4j/org/testkit/backend/RxBlockingSubscriber.java

@@ -18,23 +18,33 @@
  */
 package neo4j.org.testkit.backend;
 
-import lombok.Getter;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 
 public class RxBlockingSubscriber<T> implements Subscriber<T>
 {
-    @Getter
     private final CompletableFuture<Subscription> subscriptionFuture = new CompletableFuture<>();
+    private final CompletableFuture<Void> completionFuture = new CompletableFuture<>();
     private CompletableFuture<CompletableFuture<T>> nextSignalConsumerFuture;
 
     public void setNextSignalConsumer( CompletableFuture<T> nextSignalConsumer )
     {
         nextSignalConsumerFuture.complete( nextSignalConsumer );
     }
 
+    public CompletionStage<Subscription> getSubscriptionStage()
+    {
+        return subscriptionFuture;
+    }
+
+    public CompletionStage<Void> getCompletionStage()
+    {
+        return completionFuture;
+    }
+
     @Override
     public void onSubscribe( Subscription s )
     {
@@ -51,13 +61,13 @@ public void onNext( T t )
     @Override
     public void onError( Throwable t )
     {
-        blockUntilNextSignalConsumer().completeExceptionally( t );
+        completionFuture.completeExceptionally( t );
     }
 
     @Override
     public void onComplete()
     {
-        blockUntilNextSignalConsumer().complete( null );
+        completionFuture.complete( null );
     }
 
     private CompletableFuture<T> blockUntilNextSignalConsumer()

testkit-backend/src/main/java/neo4j/org/testkit/backend/TestkitState.java

@@ -18,53 +18,56 @@
  */
 package neo4j.org.testkit.backend;
 
-import lombok.AccessLevel;
 import lombok.Getter;
+import neo4j.org.testkit.backend.holder.AsyncSessionHolder;
+import neo4j.org.testkit.backend.holder.AsyncTransactionHolder;
+import neo4j.org.testkit.backend.holder.DriverHolder;
+import neo4j.org.testkit.backend.holder.ResultCursorHolder;
+import neo4j.org.testkit.backend.holder.ResultHolder;
+import neo4j.org.testkit.backend.holder.RxResultHolder;
+import neo4j.org.testkit.backend.holder.RxSessionHolder;
+import neo4j.org.testkit.backend.holder.RxTransactionHolder;
+import neo4j.org.testkit.backend.holder.SessionHolder;
+import neo4j.org.testkit.backend.holder.TransactionHolder;
 import neo4j.org.testkit.backend.messages.requests.TestkitCallbackResult;
 import neo4j.org.testkit.backend.messages.responses.TestkitResponse;
 import reactor.core.publisher.Mono;
 
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 
-import org.neo4j.driver.Driver;
-import org.neo4j.driver.Record;
-import org.neo4j.driver.Result;
-import org.neo4j.driver.Transaction;
-import org.neo4j.driver.async.AsyncTransaction;
-import org.neo4j.driver.async.ResultCursor;
 import org.neo4j.driver.exceptions.Neo4jException;
 import org.neo4j.driver.internal.cluster.RoutingTableRegistry;
-import org.neo4j.driver.reactive.RxResult;
-import org.neo4j.driver.reactive.RxTransaction;
 
-@Getter
 public class TestkitState
 {
+    private static final String DRIVER_NOT_FOUND_MESSAGE = "Could not find driver";
+    private static final String SESSION_NOT_FOUND_MESSAGE = "Could not find session";
     private static final String TRANSACTION_NOT_FOUND_MESSAGE = "Could not find transaction";
+    private static final String RESULT_NOT_FOUND_MESSAGE = "Could not find result";
 
-    private final Map<String,Driver> drivers = new HashMap<>();
+    private final Map<String,DriverHolder> driverIdToDriverHolder = new HashMap<>();
+    @Getter
     private final Map<String,RoutingTableRegistry> routingTableRegistry = new HashMap<>();
-    private final Map<String,SessionState> sessionStates = new HashMap<>();
-    private final Map<String,AsyncSessionState> asyncSessionStates = new HashMap<>();
-    private final Map<String,RxSessionState> rxSessionStates = new HashMap<>();
-    private final Map<String,Result> results = new HashMap<>();
-    private final Map<String,ResultCursor> resultCursors = new HashMap<>();
-    private final Map<String,RxResult> rxResults = new HashMap<>();
-    private final Map<String,RxBlockingSubscriber<Record>> rxResultIdToRecordSubscriber = new HashMap<>();
-    @Getter( AccessLevel.NONE )
-    private final Map<String,Transaction> transactions = new HashMap<>();
-    @Getter( AccessLevel.NONE )
-    private final Map<String,AsyncTransaction> asyncTransactions = new HashMap<>();
-    @Getter( AccessLevel.NONE )
-    private final Map<String,RxTransaction> rxTransactions = new HashMap<>();
+    private final Map<String,SessionHolder> sessionIdToSessionHolder = new HashMap<>();
+    private final Map<String,AsyncSessionHolder> sessionIdToAsyncSessionHolder = new HashMap<>();
+    private final Map<String,RxSessionHolder> sessionIdToRxSessionHolder = new HashMap<>();
+    private final Map<String,ResultHolder> resultIdToResultHolder = new HashMap<>();
+    private final Map<String,ResultCursorHolder> resultIdToResultCursorHolder = new HashMap<>();
+    private final Map<String,RxResultHolder> resultIdToRxResultHolder = new HashMap<>();
+    private final Map<String,TransactionHolder> transactionIdToTransactionHolder = new HashMap<>();
+    private final Map<String,AsyncTransactionHolder> transactionIdToAsyncTransactionHolder = new HashMap<>();
+    private final Map<String,RxTransactionHolder> transactionIdToRxTransactionHolder = new HashMap<>();
+    @Getter
     private final Map<String,Neo4jException> errors = new HashMap<>();
-    @Getter( AccessLevel.NONE )
     private final AtomicInteger idGenerator = new AtomicInteger( 0 );
+    @Getter
     private final Consumer<TestkitResponse> responseWriter;
+    @Getter
     private final Map<String,CompletableFuture<TestkitCallbackResult>> callbackIdToFuture = new HashMap<>();
 
     public TestkitState( Consumer<TestkitResponse> responseWriter )
@@ -77,53 +80,140 @@ public String newId()
         return String.valueOf( idGenerator.getAndIncrement() );
     }
 
-    public String addTransaction( Transaction transaction )
+    public void addDriverHolder( String id, DriverHolder driverHolder )
     {
-        String id = newId();
-        this.transactions.put( id, transaction );
-        return id;
+        driverIdToDriverHolder.put( id, driverHolder );
     }
 
-    public Transaction getTransaction( String id )
+    public DriverHolder getDriverHolder( String id )
     {
-        if ( !this.transactions.containsKey( id ) )
-        {
-            throw new RuntimeException( TRANSACTION_NOT_FOUND_MESSAGE );
-        }
-        return this.transactions.get( id );
+        return get( id, driverIdToDriverHolder, DRIVER_NOT_FOUND_MESSAGE );
     }
 
-    public String addAsyncTransaction( AsyncTransaction transaction )
+    public String addSessionHolder( SessionHolder sessionHolder )
     {
-        String id = newId();
-        this.asyncTransactions.put( id, transaction );
-        return id;
+        return add( sessionHolder, sessionIdToSessionHolder );
     }
 
-    public CompletableFuture<AsyncTransaction> getAsyncTransaction( String id )
+    public SessionHolder getSessionHolder( String id )
     {
-        if ( !this.asyncTransactions.containsKey( id ) )
-        {
-            CompletableFuture<AsyncTransaction> future = new CompletableFuture<>();
-            future.completeExceptionally( new RuntimeException( TRANSACTION_NOT_FOUND_MESSAGE ) );
-            return future;
-        }
-        return CompletableFuture.completedFuture( asyncTransactions.get( id ) );
+        return get( id, sessionIdToSessionHolder, SESSION_NOT_FOUND_MESSAGE );
+    }
+
+    public String addAsyncSessionHolder( AsyncSessionHolder sessionHolder )
+    {
+        return add( sessionHolder, sessionIdToAsyncSessionHolder );
+    }
+
+    public CompletionStage<AsyncSessionHolder> getAsyncSessionHolder( String id )
+    {
+        return getAsync( id, sessionIdToAsyncSessionHolder, SESSION_NOT_FOUND_MESSAGE );
+    }
+
+    public String addRxSessionHolder( RxSessionHolder sessionHolder )
+    {
+        return add( sessionHolder, sessionIdToRxSessionHolder );
+    }
+
+    public Mono<RxSessionHolder> getRxSessionHolder( String id )
+    {
+        return getRx( id, sessionIdToRxSessionHolder, SESSION_NOT_FOUND_MESSAGE );
+    }
+
+    public String addTransactionHolder( TransactionHolder transactionHolder )
+    {
+        return add( transactionHolder, transactionIdToTransactionHolder );
+    }
+
+    public TransactionHolder getTransactionHolder( String id )
+    {
+        return get( id, transactionIdToTransactionHolder, TRANSACTION_NOT_FOUND_MESSAGE );
+    }
+
+    public String addAsyncTransactionHolder( AsyncTransactionHolder transactionHolder )
+    {
+        return add( transactionHolder, transactionIdToAsyncTransactionHolder );
+    }
+
+    public CompletionStage<AsyncTransactionHolder> getAsyncTransactionHolder( String id )
+    {
+        return getAsync( id, transactionIdToAsyncTransactionHolder, TRANSACTION_NOT_FOUND_MESSAGE );
+    }
+
+    public String addRxTransactionHolder( RxTransactionHolder transactionHolder )
+    {
+        return add( transactionHolder, transactionIdToRxTransactionHolder );
+    }
+
+    public Mono<RxTransactionHolder> getRxTransactionHolder( String id )
+    {
+        return getRx( id, transactionIdToRxTransactionHolder, TRANSACTION_NOT_FOUND_MESSAGE );
+    }
+
+    public String addResultHolder( ResultHolder resultHolder )
+    {
+        return add( resultHolder, resultIdToResultHolder );
+    }
+
+    public ResultHolder getResultHolder( String id )
+    {
+        return get( id, resultIdToResultHolder, RESULT_NOT_FOUND_MESSAGE );
+    }
+
+    public String addAsyncResultHolder( ResultCursorHolder resultHolder )
+    {
+        return add( resultHolder, resultIdToResultCursorHolder );
+    }
+
+    public CompletionStage<ResultCursorHolder> getAsyncResultHolder( String id )
+    {
+        return getAsync( id, resultIdToResultCursorHolder, RESULT_NOT_FOUND_MESSAGE );
     }
 
-    public String addRxTransaction( RxTransaction transaction )
+    public String addRxResultHolder( RxResultHolder resultHolder )
+    {
+        return add( resultHolder, resultIdToRxResultHolder );
+    }
+
+    public Mono<RxResultHolder> getRxResultHolder( String id )
+    {
+        return getRx( id, resultIdToRxResultHolder, RESULT_NOT_FOUND_MESSAGE );
+    }
+
+    private <T> String add( T value, Map<String,T> idToT )
     {
         String id = newId();
-        this.rxTransactions.put( id, transaction );
+        idToT.put( id, value );
         return id;
     }
 
-    public Mono<RxTransaction> getRxTransaction( String id )
+    private <T> T get( String id, Map<String,T> idToT, String notFoundMessage )
+    {
+        T value = idToT.get( id );
+        if ( value == null )
+        {
+            throw new RuntimeException( notFoundMessage );
+        }
+        return value;
+    }
+
+    private <T> CompletableFuture<T> getAsync( String id, Map<String,T> idToT, String notFoundMessage )
     {
-        if ( !this.rxTransactions.containsKey( id ) )
+        CompletableFuture<T> result = new CompletableFuture<>();
+        T value = idToT.get( id );
+        if ( value == null )
+        {
+            result.completeExceptionally( new RuntimeException( notFoundMessage ) );
+        }
+        else
         {
-            return Mono.error( new RuntimeException( TRANSACTION_NOT_FOUND_MESSAGE ) );
+            result.complete( value );
         }
-        return Mono.just( rxTransactions.get( id ) );
+        return result;
+    }
+
+    private <T> Mono<T> getRx( String id, Map<String,T> idToT, String notFoundMessage )
+    {
+        return Mono.fromCompletionStage( getAsync( id, idToT, notFoundMessage ) );
     }
 }