fix wrong exception in lazy import __getattr__ (#15741)

src/prefect/__init__.py

@@ -102,17 +102,20 @@
 def __getattr__(attr_name: str) -> object:
     if attr_name in _slots:
         return _slots[attr_name]
-
-    dynamic_attr = _public_api.get(attr_name)
-    if dynamic_attr is None:
-        return importlib.import_module(f".{attr_name}", package=__name__)
-
-    package, module_name = dynamic_attr
-
-    from importlib import import_module
-
-    if module_name == "__module__":
-        return import_module(f".{attr_name}", package=package)
-    else:
-        module = import_module(module_name, package=package)
-        return getattr(module, attr_name)
+    try:
+        dynamic_attr = _public_api.get(attr_name)
+        if dynamic_attr is None:
+            return importlib.import_module(f".{attr_name}", package=__name__)
+
+        package, module_name = dynamic_attr
+
+        from importlib import import_module
+
+        if module_name == "__module__":
+            return import_module(f".{attr_name}", package=package)
+        else:
+            module = import_module(module_name, package=package)
+            return getattr(module, attr_name)
+    except ModuleNotFoundError as ex:
+        module, _, attribute = ex.name.rpartition(".")
+        raise AttributeError(f"module {module} has no attribute {attribute}") from ex

tests/test_exceptions.py

@@ -117,3 +117,23 @@ def test_pickle_roundtrip(self):
         unpickled = cloudpickle.loads(pickled)
         assert str(sme) == str(unpickled)
         assert sme.args == unpickled.args
+
+
+class TestPrefectModuleImportExceptions:
+    def test_attribute_error_on_getattr(self):
+        import prefect
+
+        with pytest.raises(
+            AttributeError, match=r"module prefect has no attribute foo"
+        ):
+            prefect.foo
+
+    def test_import_error_on_from_import(self):
+        with pytest.raises(
+            ImportError, match=r"cannot import name 'foo' from 'prefect' \(.*\)"
+        ):
+            from prefect import foo  # noqa
+
+    def test_module_not_found_erorr_on_import(self):
+        with pytest.raises(ModuleNotFoundError, match=r"No module named 'prefect.foo'"):
+            import prefect.foo  # noqa