Add Temporary:ResultKeys support in Testkit backend (#1113)

neo4j · Jan 14, 2022 · 868f846 · 868f846
1 parent 47442f5
commit 868f846
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 11 deletions.
diff --git a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java
@@ -58,7 +58,8 @@ public class GetFeatures implements TestkitRequest
             "Temporary:ConnectionAcquisitionTimeout",
             "Temporary:GetConnectionPoolMetrics",
             "Temporary:CypherPathAndRelationship",
-            "Temporary:FullSummary"
+            "Temporary:FullSummary",
+            "Temporary:ResultKeys"
     ) );
 
     private static final Set<String> SYNC_FEATURES = new HashSet<>( Arrays.asList(

diff --git a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/SessionRun.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/SessionRun.java
@@ -32,6 +32,7 @@
 import reactor.core.publisher.Mono;
 
 import java.time.Duration;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletionStage;
@@ -63,7 +64,7 @@ public TestkitResponse process( TestkitState testkitState )
         org.neo4j.driver.Result result = session.run( query, transactionConfig.build() );
         String id = testkitState.addResultHolder( new ResultHolder( sessionHolder, result ) );
 
-        return createResponse( id );
+        return createResponse( id, result.keys() );
     }
 
     @Override
@@ -86,7 +87,7 @@ public CompletionStage<TestkitResponse> processAsync( TestkitState testkitState
                                                                        {
                                                                            String id = testkitState.addAsyncResultHolder(
                                                                                    new ResultCursorHolder( sessionHolder, resultCursor ) );
-                                                                           return createResponse( id );
+                                                                           return createResponse( id, resultCursor.keys() );
                                                                        } );
                                          } );
     }
@@ -111,13 +112,13 @@ public Mono<TestkitResponse> processRx( TestkitState testkitState )
                                          // The keys() method causes RUN message exchange.
                                          // However, it does not currently report errors.
                                          return Mono.fromDirect( result.keys() )
-                                                    .map( ignored -> createResponse( id ) );
+                                                    .map( keys -> createResponse( id, keys ) );
                                      } );
     }
 
-    private Result createResponse( String resultId )
+    private Result createResponse( String resultId, List<String> keys )
     {
-        return Result.builder().data( Result.ResultBody.builder().id( resultId ).build() ).build();
+        return Result.builder().data( Result.ResultBody.builder().id( resultId ).keys( keys ).build() ).build();
     }
 
     @Setter

diff --git a/...kit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/TransactionRun.java b/...kit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/TransactionRun.java
@@ -32,6 +32,7 @@
 import reactor.core.publisher.Mono;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletionStage;
 
@@ -50,7 +51,7 @@ public TestkitResponse process( TestkitState testkitState )
         org.neo4j.driver.Result result = transactionHolder.getTransaction()
                                                           .run( data.getCypher(), data.getParams() != null ? data.getParams() : Collections.emptyMap() );
         String resultId = testkitState.addResultHolder( new ResultHolder( transactionHolder, result ) );
-        return createResponse( resultId );
+        return createResponse( resultId, result.keys() );
     }
 
     @Override
@@ -65,7 +66,7 @@ public CompletionStage<TestkitResponse> processAsync( TestkitState testkitState
                                                                                                String resultId = testkitState.addAsyncResultHolder(
                                                                                                        new ResultCursorHolder( transactionHolder,
                                                                                                                                resultCursor ) );
-                                                                                               return createResponse( resultId );
+                                                                                               return createResponse( resultId, resultCursor.keys() );
                                                                                            } ) );
     }
 
@@ -81,13 +82,13 @@ public Mono<TestkitResponse> processRx( TestkitState testkitState )
                                          String resultId = testkitState.addRxResultHolder( new RxResultHolder( transactionHolder, result ) );
                                          // The keys() method causes RUN message exchange.
                                          // However, it does not currently report errors.
-                                         return Mono.fromDirect( result.keys() ).then( Mono.just( createResponse( resultId ) ) );
+                                         return Mono.fromDirect( result.keys() ).map( keys -> createResponse( resultId, keys ) );
                                      } );
     }
 
-    protected Result createResponse( String resultId )
+    protected Result createResponse( String resultId, List<String> keys )
     {
-        return Result.builder().data( Result.ResultBody.builder().id( resultId ).build() ).build();
+        return Result.builder().data( Result.ResultBody.builder().id( resultId ).keys( keys ).build() ).build();
     }
 
     @Setter

diff --git a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/responses/Result.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/responses/Result.java
@@ -21,6 +21,8 @@
 import lombok.Builder;
 import lombok.Getter;
 
+import java.util.List;
+
 @Getter
 @Builder
 public class Result implements TestkitResponse
@@ -38,5 +40,7 @@ public String testkitName()
     public static class ResultBody
     {
         private String id;
+
+        private List<String> keys;
     }
 }