Add Result.peek support to TestKit back end (neo4j#1110) (neo4j#1133)

Co-authored-by: Robsdedude <[email protected]>

Author: injectives

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java

@@ -67,13 +67,15 @@ public class GetFeatures implements TestkitRequest
             "Feature:Bolt:3.0",
             "Optimization:PullPipelining",
             "Feature:API:Result.List",
+            "Feature:API:Result.Peek",
             "Optimization:ResultListFetchAll"
     ) );
 
     private static final Set<String> ASYNC_FEATURES = new HashSet<>( Arrays.asList(
             "Feature:Bolt:3.0",
             "Optimization:PullPipelining",
             "Feature:API:Result.List",
+            "Feature:API:Result.Peek",
             "Optimization:ResultListFetchAll"
     ) );
 

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/ResultPeek.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * 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 neo4j.org.testkit.backend.messages.requests;
+
+import lombok.Getter;
+import lombok.Setter;
+import neo4j.org.testkit.backend.RxBufferedSubscriber;
+import neo4j.org.testkit.backend.TestkitState;
+import neo4j.org.testkit.backend.holder.RxResultHolder;
+import neo4j.org.testkit.backend.messages.responses.NullRecord;
+import neo4j.org.testkit.backend.messages.responses.TestkitResponse;
+import org.neo4j.driver.Record;
+import org.neo4j.driver.Result;
+import org.neo4j.driver.exceptions.NoSuchRecordException;
+import reactor.core.publisher.Mono;
+
+import java.util.concurrent.CompletionStage;
+
+@Setter
+@Getter
+public class ResultPeek implements TestkitRequest
+{
+    private ResultPeekBody data;
+
+    @Override
+    public TestkitResponse process( TestkitState testkitState )
+    {
+        try
+        {
+            Result result = testkitState.getResultHolder( data.getResultId() ).getResult();
+            return createResponse( result.peek() );
+        }
+        catch ( NoSuchRecordException ignored )
+        {
+            return NullRecord.builder().build();
+        }
+    }
+
+    @Override
+    public CompletionStage<TestkitResponse> processAsync( TestkitState testkitState )
+    {
+        return testkitState.getAsyncResultHolder( data.getResultId() )
+                           .thenCompose( resultCursorHolder -> resultCursorHolder.getResult().peekAsync() )
+                           .thenApply( this::createResponseNullSafe );
+    }
+
+    @Override
+    public Mono<TestkitResponse> processRx( TestkitState testkitState )
+    {
+        throw new UnsupportedOperationException( "Operation not supported" );
+    }
+
+    private TestkitResponse createResponse( Record record )
+    {
+        return neo4j.org.testkit.backend.messages.responses.Record.builder()
+                                                                  .data( neo4j.org.testkit.backend.messages.responses.Record.RecordBody.builder()
+                                                                                                                                       .values( record )
+                                                                                                                                       .build() )
+                                                                  .build();
+    }
+
+    private TestkitResponse createResponseNullSafe( Record record )
+    {
+        return record != null ? createResponse( record ) : NullRecord.builder().build();
+    }
+
+    @Setter
+    @Getter
+    public static class ResultPeekBody
+    {
+        private String resultId;
+    }
+}

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/TestkitRequest.java

@@ -40,7 +40,8 @@
         @JsonSubTypes.Type( DomainNameResolutionCompleted.class ), @JsonSubTypes.Type( StartTest.class ),
         @JsonSubTypes.Type( TransactionRollback.class ), @JsonSubTypes.Type( GetFeatures.class ),
         @JsonSubTypes.Type( GetRoutingTable.class ), @JsonSubTypes.Type( TransactionClose.class ),
-        @JsonSubTypes.Type( ResultList.class ), @JsonSubTypes.Type( GetConnectionPoolMetrics.class )
+        @JsonSubTypes.Type( ResultList.class ), @JsonSubTypes.Type( GetConnectionPoolMetrics.class ),
+        @JsonSubTypes.Type( ResultPeek.class )
 } )
 public interface TestkitRequest
 {