diff --git a/services/addons/images/firmware_manager/firmware_manager.py b/services/addons/images/firmware_manager/firmware_manager.py index 04752b09f..733f92652 100644 --- a/services/addons/images/firmware_manager/firmware_manager.py +++ b/services/addons/images/firmware_manager/firmware_manager.py @@ -14,8 +14,6 @@ log_level = os.environ.get("LOGGING_LEVEL", "INFO") logging.basicConfig(format="%(levelname)s:%(message)s", level=log_level) -ACTIVE_UPGRADE_LIMIT = int(os.environ.get("ACTIVE_UPGRADE_LIMIT", "1")) - manufacturer_upgrade_scripts = { "Commsignia": "commsignia_upgrader.py", "Yunex": "yunex_upgrader.py", @@ -38,6 +36,13 @@ upgrade_queue_info = {} active_upgrades_lock = Lock() +# Changed from a constant to a function to help with unit testing +def get_upgrade_limit() -> int: + try: + upgrade_limit = int(os.environ.get("ACTIVE_UPGRADE_LIMIT", "1")) + return upgrade_limit + except ValueError: + raise ValueError("The environment variable 'ACTIVE_UPGRADE_LIMIT' must be an integer.") # Function to query the CV Manager PostgreSQL database for RSUs that have: # - A different target version than their current version @@ -71,7 +76,7 @@ def get_rsu_upgrade_data(rsu_ip="all"): def start_tasks_from_queue(): # Start the next process in the queue if there are less than ACTIVE_UPGRADE_LIMIT number of active upgrades occurring - while len(active_upgrades) < ACTIVE_UPGRADE_LIMIT and len(upgrade_queue) > 0: + while len(active_upgrades) < get_upgrade_limit() and len(upgrade_queue) > 0: rsu_to_upgrade = upgrade_queue.popleft() try: rsu_upgrade_info = upgrade_queue_info[rsu_to_upgrade] diff --git a/services/addons/tests/firmware_manager/test_firmware_manager.py b/services/addons/tests/firmware_manager/test_firmware_manager.py index 11b19b5f3..98b525797 100644 --- a/services/addons/tests/firmware_manager/test_firmware_manager.py +++ b/services/addons/tests/firmware_manager/test_firmware_manager.py @@ -2,7 +2,7 @@ from subprocess import DEVNULL from collections import deque import test_firmware_manager_values as fmv - +import pytest from addons.images.firmware_manager import firmware_manager @@ -463,7 +463,9 @@ def test_list_active_upgrades(): "addons.images.firmware_manager.firmware_manager.Popen", side_effect=Exception("Process failed to start"), ) -def test_check_for_upgrades_exception(mock_popen, mock_logging): +@patch("addons.images.firmware_manager.firmware_manager.get_upgrade_limit") +def test_check_for_upgrades_exception(mock_upgrade_limit, mock_popen, mock_logging): + mock_upgrade_limit.return_value = 5 firmware_manager.check_for_upgrades() # Assert firmware upgrade process was started with expected arguments @@ -493,7 +495,9 @@ def test_check_for_upgrades_exception(mock_popen, mock_logging): ) @patch("addons.images.firmware_manager.firmware_manager.logging") @patch("addons.images.firmware_manager.firmware_manager.start_tasks_from_queue") -def test_check_for_upgrades(mock_stfq, mock_logging): +@patch("addons.images.firmware_manager.firmware_manager.get_upgrade_limit") +def test_check_for_upgrades(mock_upgrade_limit, mock_stfq, mock_logging): + mock_upgrade_limit.return_value = 5 firmware_manager.check_for_upgrades() # Assert firmware upgrade process was started with expected arguments @@ -526,3 +530,23 @@ def test_init_background_task(mock_bgscheduler): firmware_manager.check_for_upgrades, "cron", minute="0" ) mock_bgscheduler_obj.start.assert_called_with() + + +def test_get_upgrade_limit_no_env(): + limit = firmware_manager.get_upgrade_limit() + assert limit == 1 + + +@patch.dict("os.environ", {"ACTIVE_UPGRADE_LIMIT": "5"}) +def test_get_upgrade_limit_with_env(): + limit = firmware_manager.get_upgrade_limit() + assert limit == 5 + + +@patch.dict("os.environ", {"ACTIVE_UPGRADE_LIMIT": "bad_value"}) +def test_get_upgrade_limit_with_bad_env(): + with pytest.raises( + ValueError, + match="The environment variable 'ACTIVE_UPGRADE_LIMIT' must be an integer.", + ): + firmware_manager.get_upgrade_limit() diff --git a/services/addons/tests/iss_health_check/test_iss_token.py b/services/addons/tests/iss_health_check/test_iss_token.py index 37c73eaa5..348f388f4 100644 --- a/services/addons/tests/iss_health_check/test_iss_token.py +++ b/services/addons/tests/iss_health_check/test_iss_token.py @@ -8,6 +8,7 @@ # --------------------- Storage Type tests --------------------- + @patch.dict( os.environ, { @@ -63,18 +64,14 @@ def test_get_storage_type_invalid(): iss_token.get_storage_type() -@patch.dict( - os.environ, - { - - }, -) +@patch.dict(os.environ, {}, clear=True) def test_get_storage_type_unset(): with pytest.raises(SystemExit): iss_token.get_storage_type() + # --------------------- end of Storage Type tests --------------------- - + # --------------------- GCP tests --------------------- @patch( @@ -303,20 +300,20 @@ def test_get_token_secret_exists( # Assert final value assert actual_value == expected_value + # --------------------- end of GCP tests --------------------- - + # --------------------- Postgres tests --------------------- + @patch( "addons.images.iss_health_check.iss_token.pgquery", ) def test_check_if_data_exists_true(mock_pgquery): mock_pgquery.query_db.return_value = [(1,)] actual_value = iss_token.check_if_data_exists("test-table-name") - expected_query = ( - "SELECT * FROM test-table-name" - ) + expected_query = "SELECT * FROM test-table-name" mock_pgquery.query_db.assert_called_with(expected_query) assert actual_value == True @@ -327,9 +324,7 @@ def test_check_if_data_exists_true(mock_pgquery): def test_check_if_data_exists_false(mock_pgquery): mock_pgquery.query_db.return_value = [] actual_value = iss_token.check_if_data_exists("test-table-name") - expected_query = ( - "SELECT * FROM test-table-name" - ) + expected_query = "SELECT * FROM test-table-name" mock_pgquery.query_db.assert_called_with(expected_query) assert actual_value == False @@ -352,9 +347,7 @@ def test_add_data(mock_pgquery): def test_get_latest_data(mock_pgquery): mock_pgquery.query_db.return_value = [(1, "test-common-name", "test-token")] actual_value = iss_token.get_latest_data("test-table-name") - expected_query = ( - "SELECT * FROM test-table-name ORDER BY iss_key_id DESC LIMIT 1" - ) + expected_query = "SELECT * FROM test-table-name ORDER BY iss_key_id DESC LIMIT 1" mock_pgquery.query_db.assert_called_with(expected_query) assert actual_value == {"id": 1, "name": "test-common-name", "token": "test-token"} @@ -437,7 +430,7 @@ def test_get_token_data_exists( ): # Mock every major dependency mock_check_if_data_exists.return_value = True - mock_get_latest_data.return_value = { + mock_get_latest_data.return_value = { "id": 1, "name": "test-api-key-name_01234", "token": "old-token", @@ -475,4 +468,5 @@ def test_get_token_data_exists( # Assert final value assert result == "new-iss-token" -# --------------------- end of Postgres tests --------------------- \ No newline at end of file + +# --------------------- end of Postgres tests ---------------------