Unit Test Fixes

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]

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()

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 ---------------------
+
+# --------------------- end of Postgres tests ---------------------