From 374b1a7063a3eea98be9f60258496ca0fa76bfd2 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 3 Aug 2021 16:26:50 +0200 Subject: [PATCH] UT: Add test for mongo timeout and mongo version --- monkey/monkey_island/cc/server_setup.py | 11 +++++++---- monkey/monkey_island/cc/server_utils/consts.py | 1 + .../monkey_island/cc/setup/mongo/mongo_setup.py | 14 +++++++------- .../cc/setup/mongo/test_mongo_setup.py | 16 ++++++++++++++++ 4 files changed, 31 insertions(+), 11 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 65236509d..35879a1d4 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/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( diff --git a/monkey/monkey_island/cc/server_utils/consts.py b/monkey/monkey_island/cc/server_utils/consts.py index fb5edbfab..30749cb3e 100644 --- a/monkey/monkey_island/cc/server_utils/consts.py +++ b/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)) diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index d9440abfe..196ad54bf 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/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(): - _wait_for_mongo_db_server(MONGO_URL) +def connect_to_mongodb(timeout: float): + _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: - raise MongoDBTimeOutError(f"Failed to connect to MongoDB after {TIMEOUT} seconds.") + if (time.time() - start_time) > timeout: + raise MongoDBTimeOutError(f"Failed to connect to MongoDB after {timeout} seconds.") time.sleep(1) diff --git a/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py b/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py new file mode 100644 index 000000000..502e7dbfe --- /dev/null +++ b/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)