Update implementation of import hook loaders to use find_spec.

Author: a-feld

newrelic/api/import_hook.py

@@ -18,6 +18,11 @@
 
 _logger = logging.getLogger(__name__)
 
+try:
+    from importlib.util import find_spec
+except ImportError:
+    find_spec = None
+
 _import_hooks = {}
 
 _ok_modules = (
@@ -168,14 +173,26 @@ def find_module(self, fullname, path=None):
         self._skip[fullname] = True
 
         try:
-            try:
-                from importlib.util import find_spec
-                return _ImportHookChainedLoader(find_spec(fullname).loader)
-            except AttributeError:
-                return None
-            except ImportError:
+            # For Python 3 we need to use find_spec() from the importlib
+            # module.
+
+            if find_spec:
+                spec = find_spec(fullname)
+                loader = getattr(spec, 'loader', None)
+
+                if loader:
+                    return _ImportHookChainedLoader(loader)
+
+            else:
                 __import__(fullname)
+
+                # If we get this far then the module we are
+                # interested in does actually exist and so return
+                # our loader to trigger import hooks and then return
+                # the module.
+
                 return _ImportHookLoader()
+
         finally:
             del self._skip[fullname]
 

tests/agent_unittests/test_import_hook.py

@@ -1,43 +1,47 @@
-from pytest import raises
+import sys
 
+import newrelic.api.import_hook as import_hook
+import newrelic.packages.six as six
+import pytest
 
-def test_import_hook_finder():
+# a dummy hook just to be able to register hooks for modules
+def hook(*args, **kwargs):
+    pass
+
+
+def test_import_hook_finder(monkeypatch):
     """
     This asserts the behavior of ImportHookFinder.find_module. It behaves
     differently depending on whether or not the module it is looking for
     exists, has been registered with an import hook, and across different
     python versions.
     """
-    from newrelic.api.import_hook import ImportHookFinder, register_import_hook
-    import sys
-
-    PY2 = sys.version_info[0] == 2
-
-    finder = ImportHookFinder()
+    finder = import_hook.ImportHookFinder()
 
-    # a dummy hook just to be able to register hooks for modules
-    def hook(*args, **kwargs):
-        pass
+    # Override the registered import hooks for the scope of this test
+    registered_hooks = {
+        "registered_but_does_not_exist": hook,
+        "newrelic.api": hook,
+    }
+    monkeypatch.setattr(import_hook, "_import_hooks", registered_hooks)
 
     # Finding a module that does not exist and is not registered returns None.
-    module = finder.find_module('some_module_that_does_not_exist')
+    module = finder.find_module("module_does_not_exist")
     assert module is None
 
     # Finding a module that does not exist and is registered behaves
     # differently on python 2 vs python 3.
-    register_import_hook('some_module_that_does_not_exist', hook)
-    if PY2:
-        with raises(ImportError):
-            module = finder.find_module('some_module_that_does_not_exist')
+    if six.PY2:
+        with pytest.raises(ImportError):
+            module = finder.find_module("registered_but_does_not_exist")
     else:
-        module = finder.find_module('some_module_that_does_not_exist')
+        module = finder.find_module("registered_but_does_not_exist")
         assert module is None
 
     # Finding a module that exists, but is not registered returns None.
-    module = finder.find_module('newrelic')
+    module = finder.find_module("newrelic")
     assert module is None
 
     # Finding a module that exists, and is registered, finds that module.
-    register_import_hook('newrelic', hook)
-    module = finder.find_module('newrelic')
+    module = finder.find_module("newrelic.api")
     assert module is not None