diff --git a/src/lumigo_tracer/sync_http/handler.py b/src/lumigo_tracer/sync_http/handler.py index 00e053c3..db2707cb 100644 --- a/src/lumigo_tracer/sync_http/handler.py +++ b/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 diff --git a/src/test/unit/test_handler.py b/src/test/unit/test_handler.py index 055a016c..9abd0a05 100644 --- a/src/test/unit/test_handler.py +++ b/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({}, {})