init

Signed-off-by: mgorsk1 <[email protected]>

metadata_service/proxy/gremlin_proxy.py

@@ -1735,7 +1735,8 @@ class GenericGremlinProxy(AbstractGremlinProxy):
 
     def __init__(self, *, host: str, port: Optional[int] = None, user: Optional[str] = None,
                  password: Optional[str] = None, traversal_source: 'str' = 'g', key_property_name: str = 'key',
-                 driver_remote_connection_options: Mapping[str, Any] = {}) -> None:
+                 driver_remote_connection_options: Mapping[str, Any] = {},
+                 **kwargs: dict) -> None:
         driver_remote_connection_options = dict(driver_remote_connection_options)
 
         # as others, we repurpose host a url

metadata_service/proxy/janus_graph_proxy.py

@@ -20,7 +20,8 @@ class JanusGraphGremlinProxy(AbstractGremlinProxy):
     """
     def __init__(self, *, host: str, port: Optional[int] = None, user: Optional[str] = None,
                  password: Optional[str] = None, traversal_source: 'str' = 'g',
-                 driver_remote_connection_options: Mapping[str, Any] = {}) -> None:
+                 driver_remote_connection_options: Mapping[str, Any] = {},
+                 **kwargs: dict) -> None:
         driver_remote_connection_options = dict(driver_remote_connection_options)
 
         # as others, we repurpose host a url, and url can be an HTTPRequest

metadata_service/proxy/neo4j_proxy.py

@@ -57,7 +57,8 @@ def __init__(self, *,
                  num_conns: int = 50,
                  max_connection_lifetime_sec: int = 100,
                  encrypted: bool = False,
-                 validate_ssl: bool = False) -> None:
+                 validate_ssl: bool = False,
+                 **kwargs: dict) -> None:
         """
         There's currently no request timeout from client side where server
         side can be enforced via "dbms.transaction.timeout"

metadata_service/proxy/neptune_proxy.py

@@ -57,7 +57,8 @@ class NeptuneGremlinProxy(AbstractGremlinProxy):
     def __init__(self, *, host: str, port: Optional[int] = None, user: str = None,
                  password: Optional[Union[str, boto3.session.Session]] = None,
                  driver_remote_connection_options: Mapping[str, Any] = {},
-                 neptune_bulk_loader_s3_bucket_name: Optional[str] = None) -> None:
+                 neptune_bulk_loader_s3_bucket_name: Optional[str] = None,
+                 **kwargs: dict) -> None:
 
         driver_remote_connection_options = dict(driver_remote_connection_options)
         # port should be part of that url

tests/unit/proxy/test_create_proxy.py

@@ -0,0 +1,42 @@
+# Copyright Contributors to the Amundsen project.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import unittest
+from typing import Any, Dict  # noqa: F401
+from unittest.mock import MagicMock, patch
+
+from flask import Flask
+
+import metadata_service
+from metadata_service.config import PROXY_CLIENTS
+from metadata_service.proxy import get_proxy_client
+
+
+class TestCreateProxy(unittest.TestCase):
+    proxies = {
+        'NEO4J': {'host': 'bolt://neo4j.com', 'port': 7687, 'password': 'neo4j'},
+        'ATLAS': {'host': 'http://atlas.com', 'port': 21000, 'password': 'atlas'}
+    }
+
+    @patch('neo4j.GraphDatabase.driver', MagicMock())
+    def test_proxy_client_creation(self) -> None:
+        for proxy, spec in self.proxies.items():
+            with self.subTest():
+                config = metadata_service.config.LocalConfig()
+                metadata_service.proxy._proxy_client = None
+
+                config.PROXY_CLIENT = PROXY_CLIENTS[proxy]
+                config.PROXY_HOST = spec['host']  # type: ignore
+                config.PROXY_PORT = spec['port']  # type: ignore
+
+                app = Flask(__name__)
+
+                app.config.from_object(config)
+
+                with app.app_context():
+                    try:
+                        get_proxy_client()
+                        assert True
+                    except Exception as e:
+                        assert False