Fix colo code

Author: cianyleow

states/_modules/cfutils.py

@@ -2,11 +2,10 @@
 import copy
 import logging
 import os
-
 from salt.exceptions import SaltException, ZincTimeoutError
 from salt.utils import dictupdate
 
-log = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 def dictmerge(destination: dict | None, update: dict | None, clear_none=False: bool, merge_lists=False: bool) -> dict | None:
@@ -53,36 +52,35 @@ def dictmerge_deepcopy(destination: dict | None, update: dict | None, clear_none
     return destination_copy
 
 
-def load_file_as_base64(path: str) -> bytes:
-    """
-    Read arbitrary file, and return its content as base64
-    This module exists to allow including raw config files in pillar
-    :param path: path of the file
-    :return: base64 string
-    """
-
-    if path[0] != "/":
-        path = "/etc/salt/data/" + path
-
-    with open(path, "rb") as f:
-        return base64.b64encode(f.read())
-
-
 def get_colo_names(timeout=10: int, backup=True: bool) -> list[str]:
     """
     Gets colo names using systems
-
-    :param backup (boolean)
     :param timeout (int)
+    :param backup (boolean)
     """
 
     try:
         return __salt__["zinc.get_colo_names"](timeout=(timeout * 1000))
     except ZincTimeoutError:
-        pass
+        if backup:
+            return __salt__["provision_api.get_names"](type="colo", timeout=timeout)
+        else:
+            []
     except Exception as e:
-        log.error(f"An error occurred: {e}")
+        logger.error(f"An error occurred: {e}")
         return []
 
-    if backup:
-        return __salt__["provision_api.get_names"](type="colo", timeout=timeout)
+
+def load_file_as_base64(path: str) -> bytes:
+    """
+    Read arbitrary file, and return its content as base64
+    This module exists to allow including raw config files in pillar
+    :param path: path of the file
+    :return: base64 string
+    """
+
+    if path[0] != "/":
+        path = "/etc/salt/data/" + path
+
+    with open(path, "rb") as f:
+        return base64.b64encode(f.read())

tests/pytests/unit/_modules/test_cfutils.py

@@ -1,13 +1,12 @@
-from unittest.mock import MagicMock
-
 import pytest
-
 from salt.exceptions import ProvisionAPITimeoutError
 from salt.utils import dictupdate
 from salt.testing.mocks import mock_server
 from states._modules import cfutils
 from states._modules.provision_api import get_names as provision_api_get_names
 from states._modules.zinc import get_colo_names as zinc_get_colo_names
+from unittest.mock import MagicMock
+
 
 @pytest.fixture
 def configure_loader_modules():
@@ -20,17 +19,21 @@ def configure_loader_modules():
         }
     }
 
+
 # === dictmerge tests ===
 def test_dictmerge_none_destination() -> None:
     assert cfutils.dictmerge(None, {}) == {}
 
+
 def test_dictmerge_none_update() -> None:
     assert cfutils.dictmerge({}, None) == {}
 
+
 def test_dictmerge_non_mapping_update() -> None:
     with pytest.raises(SaltException, value = "arguments must be a dictionary."):
         cfutils.dictmerge({}, "str")
 
+
 def test_dictmerge_merges_lists() -> None:
     destination = {"a": ["a"]}
     update = {"a": ["b"]}
@@ -41,20 +44,25 @@ def test_dictmerge_merges_lists() -> None:
 
     assert cfutils.dictmerge(destination, update, merge_lists = True) == expected
 
+
 def test_dictmerge_clear_none() -> None:
     assert cfutils.dictmerge({"a": None}, None) == {}
 
+
 # === dictmerge_deepcopy tests ===
 def test_dictmerge_none_destination() -> None:
     assert cfutils.dictmerge_deepcopy(None, {}) == {}
 
+
 def test_dictmerge_deepcopy_none_update() -> None:
     assert cfutils.dictmerge_deepcopy({}, None) == {}
 
+
 def test_dictmerge_deepcopy_non_mapping_update() -> None:
     with pytest.raises(SaltException, value = "arguments must be a dictionary."):
         cfutils.dictmerge_deepcopy({}, "str")
 
+
 def test_dictmerge_deepcopy_merges_lists() -> None:
     destination = {"a": ["a"]}
     update = {"a": ["b"]}
@@ -65,9 +73,11 @@ def test_dictmerge_deepcopy_merges_lists() -> None:
 
     assert cfutils.dictmerge_deepcopy(destination, update, merge_lists = True) == expected
 
+
 def test_dictmerge_deepcopy_clear_none() -> None:
     assert cfutils.dictmerge_deepcopy({"a": None}, {"b": "value"}) == {"b": "value"}
 
+
 def test_dictmerge_deepcopy_makes_copy() -> None:
     destination = {"a": "value"}
     update = {"b": "update"}
@@ -76,6 +86,7 @@ def test_dictmerge_deepcopy_makes_copy() -> None:
 
     assert cfutils.dictmerge_deepcopy(destination, update) == expected
 
+
 # === load_file_as_base64 ===
 def test_load_file_as_base64_absolute_path() -> None:
      with open("/tmp/cfutils.txt", "a") as f:
@@ -85,6 +96,7 @@ def test_load_file_as_base64_absolute_path() -> None:
 
     os.remove("/tmp/cfutils.txt")
 
+
 def test_load_file_as_base64_relative_path() -> None:
     # Create directory if it doesn't exist
     try:
@@ -105,12 +117,13 @@ def test_load_file_as_base64_relative_path() -> None:
         # Ignore errors, we definitely don't want these anymore
         pass
 
-# === get_colo_names ===
 
+# === get_colo_names ===
 def test_get_colo_names_from_zinc() -> None:
     with mock_server("zinc", "get_colo_names", 200, ["zinc_colo"]):
         assert cfutils.get_colo_names() == ["zinc_colo"]
 
+
 def test_get_colo_names_zinc_exception() -> None:
     with mock_server("zinc", "get_colo_names", 503, ""), pytest.raises(Exception, value = "An error occurred: Exception (received 503)"):
         cfutils.get_colo_names()
@@ -122,6 +135,7 @@ def test_get_colo_names_timeout_from_zinc_falls_back_to_provision_api() -> None:
     with mock_server("provision_api", "get_names", 200, ["provision_api_colo"]):
         assert cfutils.get_colo_names() == ["provision_api_colo"]
 
+
 # Really slow test due to wait for both Zinc and Provision API timeout (20s)
 @pytest.mark.slow
 def test_get_colo_names_returns_error_on_backup_failure() -> None: