Fix for zipped custom modules (#57759)

* Put '.zip' to suffix_order list as well as to suffix_map.

* Update test

* Added changelog.

* Handle .pyx modules as well + test

* Test fix: guard pyximport assertions

* Ran pre-commit hook

changelog/57674.fixed

@@ -0,0 +1 @@
+Fixed exception on loading custom zipped modules.

salt/loader.py

@@ -4,7 +4,6 @@ directories for python loadable code and organizes the code into the
 plugin interfaces used by Salt.
 """
 
-
 import functools
 import inspect
 import logging
@@ -1353,21 +1352,31 @@ class LazyLoader(salt.utils.lazy.LazyDict):
         refresh the mapping of the FS on disk
         """
         # map of suffix to description for imp
-        if self.opts.get("cython_enable", True) is True:
+        if (
+            self.opts.get("cython_enable", True) is True
+            and ".pyx" not in self.suffix_map
+        ):
             try:
                 global pyximport
                 pyximport = __import__("pyximport")  # pylint: disable=import-error
                 pyximport.install()
                 # add to suffix_map so file_mapping will pick it up
                 self.suffix_map[".pyx"] = tuple()
+                if ".pyx" not in self.suffix_order:
+                    self.suffix_order.append(".pyx")
             except ImportError:
                 log.info(
                     "Cython is enabled in the options but not present "
                     "in the system path. Skipping Cython modules."
                 )
         # Allow for zipimport of modules
-        if self.opts.get("enable_zip_modules", True) is True:
+        if (
+            self.opts.get("enable_zip_modules", True) is True
+            and ".zip" not in self.suffix_map
+        ):
             self.suffix_map[".zip"] = tuple()
+            if ".zip" not in self.suffix_order:
+                self.suffix_order.append(".zip")
         # allow for module dirs
         if USE_IMPORTLIB:
             self.suffix_map[""] = ("", "", MODULE_KIND_PKG_DIRECTORY)

tests/unit/test_loader.py

@@ -21,18 +21,12 @@ import salt.config
 import salt.loader
 import salt.utils.files
 import salt.utils.stringutils
-
-# pylint: disable=import-error,no-name-in-module,redefined-builtin
-from salt.ext import six
-from salt.ext.six.moves import range
 from tests.support.case import ModuleCase
 from tests.support.helpers import slowTest
 from tests.support.mock import MagicMock, patch
 from tests.support.runtests import RUNTIME_VARS
 from tests.support.unit import TestCase
 
-# pylint: enable=no-name-in-module,redefined-builtin
-
 log = logging.getLogger(__name__)
 
 
@@ -1595,7 +1589,7 @@ class LazyLoaderOptimizationOrderTest(TestCase):
         self.loader = self._get_loader()
         filename = self._get_module_filename()
         basename = os.path.basename(filename)
-        expected = "lazyloadertest.py" if six.PY3 else "lazyloadertest.pyc"
+        expected = "lazyloadertest.py"
         assert basename == expected, basename
 
 
@@ -1677,3 +1671,29 @@ class LazyLoaderRefreshFileMappingTest(TestCase):
         func_mock.assert_called()
         assert len(func_mock.call_args_list) == len(lock_mock.__enter__.call_args_list)
         del loader
+
+    def test_lazyloader_zip_modules(self):
+        self.opts["enable_zip_modules"] = True
+        try:
+            loader = self.__init_loader()
+            assert ".zip" in loader.suffix_map
+            assert ".zip" in loader.suffix_order
+        finally:
+            self.opts["enable_zip_modules"] = False
+        loader = self.__init_loader()
+        assert ".zip" not in loader.suffix_map
+        assert ".zip" not in loader.suffix_order
+
+    def test_lazyloader_pyx_modules(self):
+        self.opts["cython_enable"] = True
+        try:
+            loader = self.__init_loader()
+            # Don't assert if the current environment has no pyximport
+            if salt.loader.pyximport is not None:
+                assert ".pyx" in loader.suffix_map
+                assert ".pyx" in loader.suffix_order
+        finally:
+            self.opts["cython_enable"] = False
+        loader = self.__init_loader()
+        assert ".pyx" not in loader.suffix_map
+        assert ".pyx" not in loader.suffix_order