Adjust config main unit tests

aws · Sep 19, 2024 · 059e891 · 059e891
1 parent d48de20
commit 059e891
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/test/unit/configuration/test_configuration_from_main.py b/test/unit/configuration/test_configuration_from_main.py
@@ -9,7 +9,7 @@
 from graph_notebook.configuration.generate_config import AuthModeEnum, Configuration, GremlinSection
 from graph_notebook.configuration.get_config import get_config
 from graph_notebook.neptune.client import (NEPTUNE_DB_SERVICE_NAME, NEPTUNE_ANALYTICS_SERVICE_NAME,
-                                           DEFAULT_HTTP_PROTOCOL)
+                                           DEFAULT_HTTP_PROTOCOL, DEFAULT_WS_PROTOCOL)
 
 
 class TestGenerateConfigurationMain(unittest.TestCase):
@@ -135,7 +135,7 @@ def test_generate_configuration_main_gremlin_protocol_no_service(self):
         self.assertEqual(0, result)
         config = get_config(self.test_file_path)
         config_dict = config.to_dict()
-        self.assertEqual(DEFAULT_HTTP_PROTOCOL, config_dict['gremlin']['connection_protocol'])
+        self.assertEqual(DEFAULT_WS_PROTOCOL, config_dict['gremlin']['connection_protocol'])
 
     def test_generate_configuration_main_gremlin_protocol_db(self):
         result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config '
@@ -149,7 +149,7 @@ def test_generate_configuration_main_gremlin_protocol_db(self):
         self.assertEqual(0, result)
         config = get_config(self.test_file_path)
         config_dict = config.to_dict()
-        self.assertEqual(DEFAULT_HTTP_PROTOCOL, config_dict['gremlin']['connection_protocol'])
+        self.assertEqual(DEFAULT_WS_PROTOCOL, config_dict['gremlin']['connection_protocol'])
 
     def test_generate_configuration_main_gremlin_protocol_analytics(self):
         result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config '