RD-4296 - mimic AWS import strategy of the original handler (#175)

* mimic boostrap.py import strategy of the original handler

Author: saartochner-lumigo

src/lumigo_tracer/auto_instrument_handler.py

@@ -10,13 +10,13 @@ def parse_handler():
     try:
         module_name, unit_name = os.environ[ORIGINAL_HANDLER_KEY].rsplit(".", 1)
     except KeyError:
-        raise ValueError(
+        raise Exception(
             "Could not find the original handler. Please contact Lumigo for more information."
-        )
-    except ValueError:
-        raise RuntimeError(
-            f"Invalid handler format: Bad handler '{os.environ[ORIGINAL_HANDLER_KEY]}': not enough values to unpack (expected 2, got 1)"
-        )
+        ) from None
+    except ValueError as e:
+        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
 
@@ -26,11 +26,21 @@ def _handler(*args, **kwargs):
     handler_module = ""
     try:
         handler_module, unit_name = parse_handler()
-        original_handler = getattr(importlib.import_module(handler_module), unit_name)
-    except (ImportError, AttributeError):
+        original_module = importlib.import_module(handler_module)
+    except ImportError as e:
         raise ImportError(
-            f"Unable to import module '{handler_module}': No module named '{handler_module}'"
-        )
+            f"Runtime.ImportModuleError: Unable to import module '{handler_module}': {str(e)}"
+        ) from None
+    except SyntaxError as e:
+        raise SyntaxError(
+            f"Runtime.UserCodeSyntaxError: Syntax error in module '{handler_module}': {str(e)}"
+        ) from None
+    try:
+        original_handler = getattr(original_module, unit_name)
+    except AttributeError:
+        raise Exception(
+            f"Runtime.HandlerNotFound: Handler '{unit_name}' missing on module '{handler_module}'"
+        ) from None
     return original_handler(*args, **kwargs)
 
 

src/test/unit/test_auto_instrument_handler.py

@@ -32,33 +32,72 @@ def test_hierarchy_happy_flow(monkeypatch):
 def test_import_error(monkeypatch):
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "blabla.not.exists")
 
-    with pytest.raises(ImportError):
+    try:
         _handler({}, {})
+    except ImportError as e:
+        # Note: We're not using pytest.raises in order to get the exception context
+        assert "Runtime.ImportModuleError" in str(e)
+        assert "another exception occurred" not in traceback.format_exc()
+    else:
+        assert False
 
 
 def test_no_env_handler_error(monkeypatch):
-    if os.environ.get(ORIGINAL_HANDLER_KEY):
-        monkeypatch.delenv(ORIGINAL_HANDLER_KEY)
+    monkeypatch.delenv(ORIGINAL_HANDLER_KEY, None)
 
-    with pytest.raises(ValueError):
+    with pytest.raises(Exception) as e:
         _handler({}, {})
+    assert "Could not find the original handler" in str(e.value)
 
 
 def test_error_in_original_handler_no_extra_exception_log(monkeypatch, context):
-    monkeypatch.setattr(importlib, "import_module", mock.Mock(side_effect=Exception))
+    monkeypatch.setattr(importlib, "import_module", mock.Mock(side_effect=ZeroDivisionError))
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.exit")
+
+    try:
+        _handler({}, context)
+    except ZeroDivisionError:
+        # Note: We're not using pytest.raises in order to get the exception context
+        assert "another exception occurred" not in traceback.format_exc()
+    else:
+        assert False
+
+
+def test_error_in_original_handler_syntax_error(monkeypatch, context):
+    monkeypatch.setattr(importlib, "import_module", mock.Mock(side_effect=SyntaxError))
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.exit")
-    exception_occurred = False
 
     try:
         _handler({}, context)
-    except Exception:
-        exception_occurred = True
+    except SyntaxError as e:
+        # Note: We're not using pytest.raises in order to get the exception context
+        assert "Runtime.UserCodeSyntaxError" in str(e)
         assert "another exception occurred" not in traceback.format_exc()
-    assert exception_occurred is True
+    else:
+        assert False
 
 
 def test_handler_bad_format(monkeypatch):
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "no_method")
 
-    with pytest.raises(RuntimeError):
+    try:
+        _handler({}, {})
+    except ValueError as e:
+        # Note: We're not using pytest.raises in order to get the exception context
+        assert "Runtime.MalformedHandlerName" in str(e)
+        assert "another exception occurred" not in traceback.format_exc()
+    else:
+        assert False
+
+
+def test_handler_not_found(monkeypatch):
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.not_found")
+
+    try:
         _handler({}, {})
+    except Exception as e:
+        # Note: We're not using pytest.raises in order to get the exception context
+        assert "Runtime.HandlerNotFound" in str(e)
+        assert "another exception occurred" not in traceback.format_exc()
+    else:
+        assert False