change importlib to imp.load_module (#211)

Author: GuyMoses

src/lumigo_tracer/auto_instrument_handler.py

@@ -1,14 +1,20 @@
+import warnings
 import os
 import importlib
 
 from lumigo_tracer import lumigo_tracer
 
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=DeprecationWarning)
+    import imp
+
 ORIGINAL_HANDLER_KEY = "LUMIGO_ORIGINAL_HANDLER"
 
 
 def parse_handler():
     try:
         module_name, unit_name = os.environ[ORIGINAL_HANDLER_KEY].rsplit(".", 1)
+        file_handle, pathname, desc = imp.find_module(module_name)
     except KeyError:
         raise Exception(
             "Could not find the original handler. Please contact Lumigo for more information."
@@ -17,16 +23,15 @@ def parse_handler():
         raise ValueError(
             f"Runtime.MalformedHandlerName: Bad handler '{os.environ[ORIGINAL_HANDLER_KEY]}': {str(e)}"
         ) from None
-    importable_name = module_name.replace("/", ".")
-    return importable_name, unit_name
+    return module_name, unit_name, file_handle, pathname, desc
 
 
 @lumigo_tracer()
 def _handler(*args, **kwargs):
     handler_module = ""
     try:
-        handler_module, unit_name = parse_handler()
-        original_module = importlib.import_module(handler_module)
+        handler_module, unit_name, file_handle, pathname, desc = parse_handler()
+        original_module = imp.load_module(handler_module, file_handle, pathname, desc)
     except ImportError as e:
         raise ImportError(
             f"Runtime.ImportModuleError: Unable to import module '{handler_module}': {str(e)}"

src/lumigo_tracer/test_module/__init__.py

@@ -0,0 +1,2 @@
+def handler(*args, **kwargs):
+    return {"hello": "world"}

src/lumigo_tracer/test_module/test.py

@@ -1,33 +1,26 @@
-import importlib
+import imp
+
 import traceback
 
 import mock
-import sys
-import os
 
 import pytest
 from lumigo_tracer.auto_instrument_handler import _handler, ORIGINAL_HANDLER_KEY
 
 
-def test_happy_flow(monkeypatch):
-    m = mock.Mock(return_value={"hello": "world"})
-    monkeypatch.setattr(sys, "exit", m)
-    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.exit")
+def abc(*args, **kwargs):
+    return {"hello": "world"}
 
-    assert _handler({}, {}) == {"hello": "world"}
 
-    m.assert_called_once()
+def test_happy_flow(monkeypatch):
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "test_auto_instrument_handler.abc")
+    assert _handler({}, {}) == {"hello": "world"}
 
 
 def test_hierarchy_happy_flow(monkeypatch):
-    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "os/path.getsize")
-    m = mock.Mock(return_value={"hello": "world"})
-    monkeypatch.setattr(os.path, "getsize", m)
-
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "lumigo_tracer/test_module/test.handler")
     assert _handler({}, {}) == {"hello": "world"}
 
-    m.assert_called_once()
-
 
 def test_import_error(monkeypatch):
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "blabla.not.exists")
@@ -51,7 +44,7 @@ def test_no_env_handler_error(monkeypatch):
 
 
 def test_error_in_original_handler_no_extra_exception_log(monkeypatch, context):
-    monkeypatch.setattr(importlib, "import_module", mock.Mock(side_effect=ZeroDivisionError))
+    monkeypatch.setattr(imp, "load_module", mock.Mock(side_effect=ZeroDivisionError))
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.exit")
 
     try:
@@ -64,7 +57,7 @@ def test_error_in_original_handler_no_extra_exception_log(monkeypatch, context):
 
 
 def test_error_in_original_handler_syntax_error(monkeypatch, context):
-    monkeypatch.setattr(importlib, "import_module", mock.Mock(side_effect=SyntaxError))
+    monkeypatch.setattr(imp, "load_module", mock.Mock(side_effect=SyntaxError))
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.exit")
 
     try: