From 3741b801baf253d0c3c8a7aba27eed3577c400b6 Mon Sep 17 00:00:00 2001
From: Dmitriy Tverdiakov <11927660+injectives@users.noreply.github.com>
Date: Fri, 29 Oct 2021 09:43:47 +0100
Subject: [PATCH] Add Feature:Auth:Custom Testkit feature support (#1025)
 (#1051)

---
 .../messages/requests/AuthorizationToken.java | 13 +++++++++-
 .../messages/requests/GetFeatures.java        |  3 ++-
 .../backend/messages/requests/NewDriver.java  | 26 +++++++++----------
 .../backend/MessageDeserializerTest.java      |  4 +--
 4 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/AuthorizationToken.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/AuthorizationToken.java
index 3d67af0441..0e70420cdc 100644
--- a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/AuthorizationToken.java
+++ b/testkit-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
index f4b379606a..fb8fbcccd2 100644
--- 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
@@ -40,7 +40,8 @@ public class GetFeatures implements TestkitRequest
             "ConfHint:connection.recv_timeout_seconds",
             "Temporary:DriverFetchSize",
             "Temporary:DriverMaxTxRetryTime",
-            "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
index 1e01005380..6eaeddb573 100644
--- 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,24 +67,25 @@ 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 "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
index 81f309d469..3b14f34605 100644
--- 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