Add Feature:Auth:Custom Testkit feature support (#1025)

neo4j · Oct 4, 2021 · cd02074 · cd02074
1 parent a5af8e1
commit cd02074
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 18 deletions.
diff --git a/...backend/src/main/java/neo4j/org/testkit/backend/messages/requests/AuthorizationToken.java b/...backend/src/main/java/neo4j/org/testkit/backend/messages/requests/AuthorizationToken.java
@@ -31,5 +31,16 @@
 public class AuthorizationToken
 {
     @JsonProperty( "data" )
-    private Map<String,String> tokens;
+    private Tokens tokens;
+
+    @Getter
+    @Setter
+    public static class Tokens
+    {
+        private String scheme;
+        private String principal;
+        private String credentials;
+        private String realm;
+        private Map<String,Object> parameters;
+    }
 }
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
@@ -41,7 +41,8 @@ public class GetFeatures implements TestkitRequest
             "Temporary:DriverFetchSize",
             "Temporary:DriverMaxTxRetryTime",
             "Feature:Auth:Bearer",
-            "Feature:Auth:Kerberos"
+            "Feature:Auth:Kerberos",
+            "Feature:Auth:Custom"
     ) );
 
     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/NewDriver.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/NewDriver.java
@@ -23,7 +23,6 @@
 import lombok.Setter;
 import neo4j.org.testkit.backend.TestkitState;
 import neo4j.org.testkit.backend.holder.DriverHolder;
-import neo4j.org.testkit.backend.messages.responses.BackendError;
 import neo4j.org.testkit.backend.messages.responses.DomainNameResolutionRequired;
 import neo4j.org.testkit.backend.messages.responses.Driver;
 import neo4j.org.testkit.backend.messages.responses.DriverError;
@@ -68,27 +67,28 @@ public TestkitResponse process( TestkitState testkitState )
         String id = testkitState.newId();
 
         AuthToken authToken;
-        switch ( data.getAuthorizationToken().getTokens().get( "scheme" ) )
+        switch ( data.getAuthorizationToken().getTokens().getScheme() )
         {
         case "basic":
-            authToken = AuthTokens.basic( data.authorizationToken.getTokens().get( "principal" ),
-                                          data.authorizationToken.getTokens().get( "credentials" ),
-                                          data.authorizationToken.getTokens().get( "realm" ) );
+            authToken = AuthTokens.basic( data.authorizationToken.getTokens().getPrincipal(),
+                                          data.authorizationToken.getTokens().getCredentials(),
+                                          data.authorizationToken.getTokens().getRealm() );
             break;
         case "bearer":
-            authToken = AuthTokens.bearer( data.authorizationToken.getTokens().get( "credentials" ) );
+            authToken = AuthTokens.bearer( data.authorizationToken.getTokens().getCredentials() );
             break;
         case "kerberos":
-            authToken = AuthTokens.kerberos( data.authorizationToken.getTokens().get( "credentials" ) );
+            authToken = AuthTokens.kerberos( data.authorizationToken.getTokens().getCredentials() );
             break;
         default:
-            return BackendError.builder()
-                               .data( BackendError
-                                              .BackendErrorBody.builder()
-                                                               .msg( "Auth scheme " + data.authorizationToken.getTokens().get( "scheme" ) +
-                                                                     " not implemented" )
-                                                               .build() )
-                               .build();
+            authToken = AuthTokens.custom(
+                    data.authorizationToken.getTokens().getPrincipal(),
+                    data.authorizationToken.getTokens().getCredentials(),
+                    data.authorizationToken.getTokens().getRealm(),
+                    data.authorizationToken.getTokens().getScheme(),
+                    data.authorizationToken.getTokens().getParameters()
+            );
+            break;
         }
 
         Config.ConfigBuilder configBuilder = Config.builder();

diff --git a/testkit-backend/src/test/java/neo4j/org/testkit/backend/MessageDeserializerTest.java b/testkit-backend/src/test/java/neo4j/org/testkit/backend/MessageDeserializerTest.java
@@ -48,8 +48,8 @@ void testDeserializeNewDriver() throws JsonProcessingException
 
         NewDriver newDriver = (NewDriver) message;
         assertThat( newDriver.getData().getUri(), equalTo( "bolt://localhost:7687" ) );
-        assertThat( newDriver.getData().getAuthorizationToken().getTokens().get( "scheme" ), equalTo( "basic" ) );
-        assertThat( newDriver.getData().getAuthorizationToken().getTokens().get( "principal" ), equalTo( "neo4j" ) );
+        assertThat( newDriver.getData().getAuthorizationToken().getTokens().getScheme(), equalTo( "basic" ) );
+        assertThat( newDriver.getData().getAuthorizationToken().getTokens().getPrincipal(), equalTo( "neo4j" ) );
     }
 
     @Test