auto-instrument: import the user handler during import (#142)

Author: saartochner-lumigo

src/lumigo_tracer/sync_http/handler.py

@@ -6,18 +6,36 @@
 ORIGINAL_HANDLER_KEY = "LUMIGO_ORIGINAL_HANDLER"
 
 
+def parse_handler():
+    try:
+        module_name, unit_name = os.environ[ORIGINAL_HANDLER_KEY].rsplit(".", 1)
+    except KeyError:
+        raise ValueError(
+            "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)"
+        )
+    importable_name = module_name.replace("/", ".")
+    return importable_name, unit_name
+
+
 @lumigo_tracer()
 def _handler(*args, **kwargs):
+    handler_module = ""
     try:
-        module_name, unit_name = os.environ[ORIGINAL_HANDLER_KEY].rsplit(".", 1)
-        importable_name = module_name.replace("/", ".")
-        original_handler = getattr(importlib.import_module(importable_name), unit_name)
+        handler_module, unit_name = parse_handler()
+        original_handler = getattr(importlib.import_module(handler_module), unit_name)
     except (ImportError, AttributeError):
         raise ImportError(
-            "Could not load the original handler. Are you sure that the import is ok?"
-        )
-    except KeyError:
-        raise ValueError(
-            "Could not find the original handler. Please follow lumigo's docs: https://docs.lumigo.io/"
+            f"Unable to import module '{handler_module}': No module named '{handler_module}'"
         )
     return original_handler(*args, **kwargs)
+
+
+try:
+    # import handler during runtime initialization, as usual.
+    importlib.import_module(parse_handler()[0])
+except Exception:
+    pass

src/test/unit/test_handler.py

@@ -55,3 +55,10 @@ def test_error_in_original_handler_no_extra_exception_log(monkeypatch, context):
         exception_occurred = True
         assert "another exception occurred" not in traceback.format_exc()
     assert exception_occurred is True
+
+
+def test_handler_bad_format(monkeypatch):
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "no_method")
+
+    with pytest.raises(RuntimeError):
+        _handler({}, {})