Island: Extract start and connect mongodb functions

monkey/monkey_island/cc/server_setup.py

@@ -32,15 +32,9 @@ from monkey_island.cc.services.utils.network_utils import local_ip_addresses  #
 from monkey_island.cc.setup import island_config_options_validator  # noqa: E402
 from monkey_island.cc.setup.gevent_hub_error_handler import GeventHubErrorHandler  # noqa: E402
 from monkey_island.cc.setup.island_config_options import IslandConfigOptions  # noqa: E402
+from monkey_island.cc.setup.mongo import mongo_setup  # noqa: E402
 from monkey_island.cc.setup.mongo.database_initializer import init_collections  # noqa: E402
-from monkey_island.cc.setup.mongo.mongo_setup import (  # noqa: E402
+from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcess  # noqa: E402
-    MONGO_URL,
-    MongoDBTimeOutException,
-    MongoDBVersionException,
-    connect_to_mongodb,
-    register_mongo_shutdown_callback,
-    start_mongodb,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -55,24 +49,10 @@ def run_monkey_island():
     _initialize_globals(config_options, server_config_path)
 
     mongo_db_process = None
-
     if config_options.start_mongodb:
-        mongo_db_process = start_mongodb(config_options.data_dir)
+        mongo_db_process = _start_mongodb(config_options.data_dir)
-        register_mongo_shutdown_callback(mongo_db_process)
 
-    try:
+    _connect_to_mongodb(config_options.data_dir, mongo_db_process)
-        connect_to_mongodb()
-    except MongoDBTimeOutException as ex:
-        if config_options.start_mongodb and not mongo_db_process.is_running():
-            logger.error(
-                f"Failed to start MongoDB process. Check log at {mongo_db_process.log_file}."
-            )
-        else:
-            logger.error(ex)
-        sys.exit(1)
-    except MongoDBVersionException as ex:
-        logger.error(ex)
-        sys.exit(1)
 
     _configure_gevent_exception_handling(Path(config_options.data_dir))
     _start_island_server(island_args.setup_only, config_options)
@@ -109,6 +89,29 @@ def _initialize_globals(config_options: IslandConfigOptions, server_config_path:
     initialize_services(config_options.data_dir)
 
 
+def _start_mongodb(data_dir: Path) -> MongoDbProcess:
+    mongo_db_process = mongo_setup.start_mongodb(data_dir)
+    mongo_setup.register_mongo_shutdown_callback(mongo_db_process)
+
+    return mongo_db_process
+
+
+def _connect_to_mongodb(data_dir: Path, mongo_db_process: MongoDbProcess):
+    try:
+        mongo_setup.connect_to_mongodb()
+    except mongo_setup.MongoDBTimeOutException as ex:
+        if mongo_db_process and not mongo_db_process.is_running():
+            logger.error(
+                f"Failed to start MongoDB process. Check log at {mongo_db_process.log_file}."
+            )
+        else:
+            logger.error(ex)
+        sys.exit(1)
+    except mongo_setup.MongoDBVersionException as ex:
+        logger.error(ex)
+        sys.exit(1)
+
+
 def _configure_gevent_exception_handling(data_dir):
     hub = gevent.hub.get_hub()
 
@@ -123,7 +126,7 @@ def _configure_gevent_exception_handling(data_dir):
 
 def _start_island_server(should_setup_only, config_options: IslandConfigOptions):
     populate_exporter_list()
-    app = init_app(MONGO_URL)
+    app = init_app(mongo_setup.MONGO_URL)
 
     init_collections()