UT: Add test for mongo timeout and mongo version

monkey/monkey_island/cc/server_setup.py

@@ -23,7 +23,10 @@ from monkey_island.cc.arg_parser import IslandCmdArgs  # noqa: E402
 from monkey_island.cc.arg_parser import parse_cli_args  # noqa: E402
 from monkey_island.cc.resources.monkey_download import MonkeyDownload  # noqa: E402
 from monkey_island.cc.server_utils.bootloader_server import BootloaderHttpServer  # noqa: E402
-from monkey_island.cc.server_utils.consts import GEVENT_EXCEPTION_LOG  # noqa: E402
+from monkey_island.cc.server_utils.consts import (  # noqa: E402
+    GEVENT_EXCEPTION_LOG,
+    MONGO_CONNECTION_TIMEOUT,
+)
 from monkey_island.cc.server_utils.encryptor import initialize_encryptor  # noqa: E402
 from monkey_island.cc.server_utils.island_logger import reset_logger, setup_logging  # noqa: E402
 from monkey_island.cc.services.initialize import initialize_services  # noqa: E402
@@ -52,7 +55,7 @@ def run_monkey_island():
     if config_options.start_mongodb:
         mongo_db_process = _start_mongodb(config_options.data_dir)
 
-    _connect_to_mongodb(config_options.data_dir, mongo_db_process)
+    _connect_to_mongodb(mongo_db_process)
 
     _configure_gevent_exception_handling(Path(config_options.data_dir))
     _start_island_server(island_args.setup_only, config_options)
@@ -96,9 +99,9 @@ def _start_mongodb(data_dir: Path) -> MongoDbProcess:
     return mongo_db_process
 
 
-def _connect_to_mongodb(data_dir: Path, mongo_db_process: MongoDbProcess):
+def _connect_to_mongodb(mongo_db_process: MongoDbProcess):
     try:
-        mongo_setup.connect_to_mongodb()
+        mongo_setup.connect_to_mongodb(MONGO_CONNECTION_TIMEOUT)
     except mongo_setup.MongoDBTimeOutError as ex:
         if mongo_db_process and not mongo_db_process.is_running():
             logger.error(

monkey/monkey_island/cc/server_utils/consts.py

@@ -34,6 +34,7 @@ _MONGO_EXECUTABLE_PATH_LINUX = os.path.join(_MONGO_BINARY_DIR, "bin", "mongod")
 MONGO_EXECUTABLE_PATH = (
     _MONGO_EXECUTABLE_PATH_WIN if is_windows_os() else _MONGO_EXECUTABLE_PATH_LINUX
 )
+MONGO_CONNECTION_TIMEOUT = 15
 
 DEFAULT_SERVER_CONFIG_PATH = str(Path(MONKEY_ISLAND_ABS_PATH, "cc", SERVER_CONFIG_FILENAME))
 

monkey/monkey_island/cc/setup/mongo/mongo_setup.py

@@ -2,6 +2,7 @@ import atexit
 import logging
 import os
 import time
+from pathlib import Path
 
 from monkey_island.cc.database import get_db_version, is_db_server_up
 from monkey_island.cc.server_utils.file_utils import create_secure_directory
@@ -16,12 +17,11 @@ MONGO_URL = os.environ.get(
     "mongodb://{0}:{1}/{2}".format(MONGO_DB_HOST, MONGO_DB_PORT, MONGO_DB_NAME),
 )
 MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0"
-TIMEOUT = 15
 
 logger = logging.getLogger(__name__)
 
 
-def start_mongodb(data_dir: str) -> MongoDbProcess:
+def start_mongodb(data_dir: Path) -> MongoDbProcess:
     db_dir = _create_db_dir(data_dir)
     log_file = os.path.join(data_dir, MONGO_LOG_FILENAME)
 
@@ -43,20 +43,20 @@ def register_mongo_shutdown_callback(mongo_db_process: MongoDbProcess):
     atexit.register(mongo_db_process.stop)
 
 
-def connect_to_mongodb():
+def connect_to_mongodb(timeout: float):
-    _wait_for_mongo_db_server(MONGO_URL)
+    _wait_for_mongo_db_server(MONGO_URL, timeout)
     _assert_mongo_db_version(MONGO_URL)
     mongo_connector.connect_dal_to_mongodb()
 
 
-def _wait_for_mongo_db_server(mongo_url):
+def _wait_for_mongo_db_server(mongo_url, timeout):
     start_time = time.time()
 
     while not is_db_server_up(mongo_url):
         logger.info(f"Waiting for MongoDB server on {mongo_url}")
 
-        if (time.time() - start_time) > TIMEOUT:
+        if (time.time() - start_time) > timeout:
-            raise MongoDBTimeOutError(f"Failed to connect to MongoDB after {TIMEOUT} seconds.")
+            raise MongoDBTimeOutError(f"Failed to connect to MongoDB after {timeout} seconds.")
 
         time.sleep(1)
 

monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py

@@ -0,0 +1,16 @@
+import pytest
+
+from monkey_island.cc.setup.mongo import mongo_setup
+
+
+def test_connect_to_mongodb_timeout(monkeypatch):
+    monkeypatch.setattr(mongo_setup, "is_db_server_up", lambda _: False)
+    with pytest.raises(mongo_setup.MongoDBTimeOutError):
+        mongo_setup.connect_to_mongodb(0.0000000001)
+
+
+def test_connect_to_mongodb_version_too_old(monkeypatch):
+    monkeypatch.setattr(mongo_setup, "is_db_server_up", lambda _: True)
+    monkeypatch.setattr(mongo_setup, "get_db_version", lambda _: ("1", "0", "0"))
+    with pytest.raises(mongo_setup.MongoDBVersionError):
+        mongo_setup.connect_to_mongodb(0)