From 922c456977f3e18022b549154470d1faf3bbecc8 Mon Sep 17 00:00:00 2001
From: Ben Bariteau <bariteau@yelp.com>
Date: Tue, 3 Oct 2023 14:05:23 -0700
Subject: [PATCH] replace usages of imp with importlib

---
 test/discovery_failure_test.py |  6 ++++++
 test/test_runner_test.py       |  4 ++--
 testify/test_program.py        | 13 +++----------
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/test/discovery_failure_test.py b/test/discovery_failure_test.py
index 638c881..39bfbe3 100644
--- a/test/discovery_failure_test.py
+++ b/test/discovery_failure_test.py
@@ -35,8 +35,11 @@ def test_discover_test_with_broken_import(self):
                 r'    Traceback \(most recent call last\):\n'
                 r'      File "[^"]+", line \d+, in discover\n'
                 r'        submod = __import__\(module_name, fromlist=\[str\(\'__trash\'\)\]\)\n'
+                r'(                 \^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^'
+                r'\^\^\^\^\n)?'
                 r'      File "[^"]+", line \d+, in <module>\n'
                 r'        import non_existent_module  \# noqa: F401\n'
+                r'(        \^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\n)?'
                 r'    (ModuleNotFoundError|ImportError): No module named \'?non_existent_module\'?\n'
             ),
         )
@@ -51,8 +54,11 @@ def test_discover_test_with_broken_import(self):
                 r'Traceback \(most recent call last\):\n'
                 r'  File .+, line \d+, in discover\n'
                 r"    submod = __import__\(module_name, fromlist=\[str\(\'__trash\'\)\]\)\n"
+                r'(             \^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^'
+                r'\^\^\n)?'
                 r'  File .+, line \d+, in <module>\n'
                 r'    import non_existent_module  \# noqa: F401\n'
+                r'(    \^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\^\n)?'
                 r"(ModuleNotFoundError|ImportError): No module named '?non_existent_module'?\n"
             ),
         )
diff --git a/test/test_runner_test.py b/test/test_runner_test.py
index ed24d45..fbf106e 100644
--- a/test/test_runner_test.py
+++ b/test/test_runner_test.py
@@ -1,10 +1,10 @@
-import imp
 import mock
 from testify import assert_equal
 from testify import setup
 from testify import setup_teardown
 from testify import test_case
 from testify import test_runner
+from types import ModuleType
 
 from .test_runner_subdir.inheriting_class import InheritingClass
 
@@ -59,7 +59,7 @@ class PluginTestCase(test_case.TestCase):
     """
     @setup
     def build_module(self):
-        self.our_module = imp.new_module("our_module")
+        self.our_module = ModuleType("our_module")
         setattr(self.our_module, "prepare_test_case", prepare_test_case)
         setattr(self.our_module, "run_test_case", run_test_case)
         setattr(self.our_module, "add_testcase_info", add_testcase_info)
diff --git a/testify/test_program.py b/testify/test_program.py
index 20f435f..606908e 100644
--- a/testify/test_program.py
+++ b/testify/test_program.py
@@ -19,7 +19,7 @@
 import pprint
 import sys
 import logging
-import imp
+from importlib.machinery import SourceFileLoader
 
 import testify
 from testify import exit
@@ -46,11 +46,6 @@ def load_plugins():
     # The idea will be to check out the directory contents and pick up any files that seem to match what python knows how to
     # import.
 
-    # To properly load the module, we'll need to identify what type it is by the file extension
-    suffix_map = {}
-    for suffix in imp.get_suffixes():
-        suffix_map[suffix[0]] = suffix
-
     plugin_directories = [DEFAULT_PLUGIN_PATH]
     if 'TESTIFY_PLUGIN_PATH' in os.environ:
         plugin_directories += os.environ['TESTIFY_PLUGIN_PATH'].split(':')
@@ -60,16 +55,14 @@ def load_plugins():
         for file_name in os.listdir(plugin_path):
 
             # For any file that we know how to load, try to import it
-            if any(file_name.endswith('.py') and not file_name.startswith('.') for suffix in suffix_map.keys()):
+            if file_name.endswith('.py') and not file_name.startswith('.'):
                 full_file_path = os.path.join(plugin_path, file_name)
                 mod_name, suffix = os.path.splitext(file_name)
                 # Need some unlikely-to-clash unique-ish module name
                 mod_name = '_testify_plugin__' + mod_name
 
                 try:
-                    plugin_modules.append(
-                        imp.load_source(mod_name, full_file_path),
-                    )
+                    SourceFileLoader(mod_name, full_file_path).load_module()
                 except TypeError:
                     continue
                 except ImportError as e: