Add test

test/auth_module/dummy_auth_module.py

@@ -0,0 +1,21 @@
+#!/usr/bin/python3
+import io
+import json
+
+
+def authenticate(scheme: str, response: str):
+    return {
+        "authenticated": True,
+        "role": "architect",
+        "username": "andy",
+    }
+
+
+if __name__ == "__main__":
+    # I/O with Memgraph
+    input_stream = io.FileIO(1000, mode="r")
+    output_stream = io.FileIO(1001, mode="w")
+    while True:
+        params = json.loads(input_stream.readline().decode("ascii"))
+        ret = authenticate(**params)
+        output_stream.write((json.dumps(ret) + "\n").encode("ascii"))

test/common.py

@@ -70,7 +70,7 @@ def terminate(self):
             self.process.wait()
 
 
-def start_memgraph(cert_file="", key_file=""):
+def start_memgraph(cert_file="", key_file="", auth_module_mappings=""):
     if MEMGRAPH_HOST:
         use_ssl = MEMGRAPH_STARTED_WITH_SSL is not None
         return Memgraph(MEMGRAPH_HOST, MEMGRAPH_PORT, use_ssl, None)
@@ -94,6 +94,8 @@ def start_memgraph(cert_file="", key_file=""):
         "--log-file",
         "",
     ]
+    if auth_module_mappings:
+        cmd.insert(-2, f"--auth-module-mappings={auth_module_mappings}")
     memgraph_process = subprocess.Popen(cmd)
     wait_for_server(MEMGRAPH_PORT)
     use_ssl = True if key_file.strip() else False

test/test_connection.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import tempfile
+
 import mgclient
 import pytest
-import tempfile
 
 from common import start_memgraph, Memgraph, requires_ssl_enabled, requires_ssl_disabled
 from OpenSSL import crypto
@@ -63,6 +65,42 @@ def secure_memgraph_server():
     memgraph.terminate()
 
 
+@pytest.fixture(scope="function")
+def provide_role():
+    memgraph = start_memgraph()
+    conn = mgclient.connect(
+        host=memgraph.host,
+        port=memgraph.port,
+    )
+    conn.autocommit = True
+    cursor = conn.cursor()
+    cursor.execute("CREATE ROLE architect;")
+    memgraph.terminate()
+
+    yield None
+
+    memgraph = start_memgraph()
+    conn = mgclient.connect(
+        host=memgraph.host,
+        port=memgraph.port,
+    )
+    conn.autocommit = True
+    cursor = conn.cursor()
+    cursor.execute("DROP ROLE architect;")
+    memgraph.terminate()
+
+
+@pytest.fixture(scope="function")
+def auth_module_path():
+    yield os.path.normpath(
+        os.path.join(
+            os.path.dirname(os.path.realpath(__file__)),
+            "auth_module",
+            "dummy_auth_module.py",
+        )
+    )
+
+
 def test_connect_args_validation():
     # bad port
     with pytest.raises(ValueError):
@@ -82,6 +120,22 @@ def test_connect_args_validation():
         )
 
 
+def test_connect_with_custom_auth_scheme(provide_role, auth_module_path):
+    custom_scheme = "custom_scheme"
+    memgraph = start_memgraph(
+        auth_module_mappings=f"{custom_scheme}:{auth_module_path}"
+    )
+    conn = mgclient.connect(
+        host=memgraph.host,
+        port=memgraph.port,
+        scheme=custom_scheme,
+        username="andy",
+        password="dummy auth token",
+    )
+    assert conn.status == mgclient.CONN_STATUS_READY
+    memgraph.terminate()
+
+
 @requires_ssl_disabled
 def test_connect_insecure_success(memgraph_server):
     host, port, sslmode, _ = memgraph_server