Skip to content

Commit f669784

Browse files
authored
Merge pull request #142 from evakhoni/import/python-pr-805
Gracefully handle missing driver client imports
2 parents a3d3bae + 085af1f commit f669784

8 files changed

Lines changed: 222 additions & 36 deletions

File tree

python/packages/jumpstarter/jumpstarter/client/base.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from pydantic.dataclasses import dataclass
1313

1414
from .core import AsyncDriverClient
15+
from jumpstarter.common.importlib import _format_missing_driver_message
1516
from jumpstarter.streams.blocking import BlockingStream
1617

1718

@@ -103,3 +104,54 @@ def close(self):
103104

104105
def __del__(self):
105106
self.close()
107+
108+
109+
@dataclass(kw_only=True, config=ConfigDict(arbitrary_types_allowed=True))
110+
class StubDriverClient(DriverClient):
111+
"""Stub client for drivers that are not installed.
112+
113+
This client is created when a driver client class cannot be imported.
114+
It provides a placeholder that raises a clear error when the driver
115+
is actually used.
116+
"""
117+
118+
def _get_missing_class_path(self) -> str:
119+
"""Get the missing class path from labels."""
120+
return self.labels["jumpstarter.dev/client"]
121+
122+
def _raise_missing_error(self):
123+
"""Raise ImportError with installation instructions."""
124+
class_path = self._get_missing_class_path()
125+
message = _format_missing_driver_message(class_path)
126+
raise ImportError(message)
127+
128+
def call(self, method, *args):
129+
"""Invoke driver call - raises ImportError since driver is not installed."""
130+
self._raise_missing_error()
131+
132+
def streamingcall(self, method, *args):
133+
"""Invoke streaming driver call - raises ImportError since driver is not installed."""
134+
self._raise_missing_error()
135+
# Unreachable yield to make this a generator function for type checking
136+
while False: # noqa: SIM114
137+
yield
138+
139+
@contextmanager
140+
def stream(self, method="connect"):
141+
"""Open a stream - raises ImportError since driver is not installed."""
142+
self._raise_missing_error()
143+
yield
144+
145+
@contextmanager
146+
def log_stream(self):
147+
"""Open a log stream - raises ImportError since driver is not installed."""
148+
self._raise_missing_error()
149+
yield
150+
151+
def __getattr__(self, name):
152+
"""Catch any attribute access and raise the missing driver error.
153+
154+
This ensures that calls like .on(), .off(), .write() etc. on stub clients
155+
raise a helpful ImportError instead of AttributeError.
156+
"""
157+
self._raise_missing_error()
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
"""Tests for StubDriverClient."""
2+
3+
import logging
4+
from contextlib import ExitStack
5+
from unittest.mock import MagicMock, create_autospec
6+
from uuid import uuid4
7+
8+
import pytest
9+
from anyio.from_thread import BlockingPortal
10+
11+
from .base import StubDriverClient
12+
from jumpstarter.common.utils import serve
13+
from jumpstarter.driver import Driver
14+
15+
16+
class MissingClientDriver(Driver):
17+
"""Test driver that returns a non-existent client class path."""
18+
19+
@classmethod
20+
def client(cls) -> str:
21+
return "nonexistent_driver_package.client.NonExistentClient"
22+
23+
24+
def create_stub_client(class_path: str) -> StubDriverClient:
25+
"""Create a StubDriverClient with minimal mocking for testing."""
26+
return StubDriverClient(
27+
uuid=uuid4(),
28+
labels={"jumpstarter.dev/client": class_path},
29+
stub=MagicMock(),
30+
portal=create_autospec(BlockingPortal, instance=True),
31+
stack=ExitStack(),
32+
)
33+
34+
35+
def test_missing_driver_logs_warning_and_creates_stub(caplog):
36+
"""Test that a missing driver logs a warning and creates a StubDriverClient."""
37+
expected_class_path = "nonexistent_driver_package.client.NonExistentClient"
38+
with caplog.at_level(logging.WARNING):
39+
with serve(MissingClientDriver()) as client:
40+
# Should have logged a warning with the exact class path from MissingDriverError
41+
assert f"Driver client '{expected_class_path}' is not available." in caplog.text
42+
43+
# Should have created a StubDriverClient
44+
assert isinstance(client, StubDriverClient)
45+
46+
# Using the stub should raise an error
47+
with pytest.raises(ImportError):
48+
client.call("some_method")
49+
50+
51+
def test_stub_driver_client_streamingcall_raises():
52+
"""Test that streamingcall() raises ImportError with driver info."""
53+
stub = create_stub_client("missing_driver.client.Client")
54+
with pytest.raises(ImportError) as exc_info:
55+
# Need to consume the generator to trigger the error
56+
list(stub.streamingcall("some_method"))
57+
assert "missing_driver" in str(exc_info.value)
58+
59+
60+
def test_stub_driver_client_stream_raises():
61+
"""Test that stream() raises ImportError with driver info."""
62+
stub = create_stub_client("missing_driver.client.Client")
63+
with pytest.raises(ImportError) as exc_info:
64+
with stub.stream():
65+
pass
66+
assert "missing_driver" in str(exc_info.value)
67+
68+
69+
def test_stub_driver_client_log_stream_raises():
70+
"""Test that log_stream() raises ImportError with driver info."""
71+
stub = create_stub_client("missing_driver.client.Client")
72+
with pytest.raises(ImportError) as exc_info:
73+
with stub.log_stream():
74+
pass
75+
assert "missing_driver" in str(exc_info.value)
76+
77+
78+
def test_stub_driver_client_error_message_jumpstarter_driver():
79+
"""Test that error message mentions version mismatch for Jumpstarter drivers."""
80+
stub = create_stub_client("jumpstarter_driver_xyz.client.XyzClient")
81+
with pytest.raises(ImportError) as exc_info:
82+
stub.call("some_method")
83+
assert "version mismatch" in str(exc_info.value)
84+
85+
86+
def test_stub_driver_client_error_message_third_party():
87+
"""Test that error message includes install instructions for third-party drivers."""
88+
stub = create_stub_client("custom_driver.client.CustomClient")
89+
with pytest.raises(ImportError) as exc_info:
90+
stub.call("some_method")
91+
assert "pip install custom_driver" in str(exc_info.value)
92+
93+
94+
def test_stub_driver_client_arbitrary_method_raises():
95+
"""Test that accessing arbitrary methods like .on() raises ImportError, not AttributeError."""
96+
stub = create_stub_client("jumpstarter_driver_power.client.PowerClient")
97+
# Accessing .on() should raise ImportError with helpful message, not AttributeError
98+
with pytest.raises(ImportError) as exc_info:
99+
stub.on()
100+
assert "jumpstarter_driver_power" in str(exc_info.value)
101+
assert "version mismatch" in str(exc_info.value)

python/packages/jumpstarter/jumpstarter/client/client.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import logging
2+
import os
13
from collections import OrderedDict, defaultdict
24
from contextlib import ExitStack, asynccontextmanager
35
from graphlib import TopologicalSorter
@@ -9,8 +11,12 @@
911

1012
from .grpc import MultipathExporterStub
1113
from jumpstarter.client import DriverClient
14+
from jumpstarter.client.base import StubDriverClient
15+
from jumpstarter.common.exceptions import MissingDriverError
1216
from jumpstarter.common.importlib import import_class
1317

18+
logger = logging.getLogger(__name__)
19+
1420

1521
@asynccontextmanager
1622
async def client_from_path(path: str, portal: BlockingPortal, stack: ExitStack, allow: list[str], unsafe: bool):
@@ -50,7 +56,14 @@ async def client_from_channel(
5056
for index in TopologicalSorter(topo).static_order():
5157
report = reports[index]
5258

53-
client_class = import_class(report.labels["jumpstarter.dev/client"], allow, unsafe)
59+
try:
60+
client_class = import_class(report.labels["jumpstarter.dev/client"], allow, unsafe)
61+
except MissingDriverError as e:
62+
# Create stub client instead of failing
63+
# Suppress duplicate warnings
64+
if not os.environ.get("_JMP_SUPPRESS_DRIVER_WARNINGS"):
65+
logger.warning("Driver client '%s' is not available.", e.class_path)
66+
client_class = StubDriverClient
5467

5568
client = client_class(
5669
uuid=UUID(report.uuid),

python/packages/jumpstarter/jumpstarter/common/exceptions.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,3 +79,15 @@ class EnvironmentVariableNotSetError(JumpstarterException):
7979
"""Raised when a environment variable is not set."""
8080

8181
pass
82+
83+
84+
class MissingDriverError(JumpstarterException):
85+
"""Raised when a driver module is not found but should be handled gracefully.
86+
87+
This exception is raised when a driver client class cannot be imported,
88+
but the connection should continue with a stub client instead of failing.
89+
"""
90+
91+
def __init__(self, message: str, class_path: str):
92+
super().__init__(message)
93+
self.class_path = class_path

python/packages/jumpstarter/jumpstarter/common/importlib.py

Lines changed: 33 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,33 @@
55
from fnmatch import fnmatchcase
66
from importlib import import_module
77

8+
from jumpstarter.common.exceptions import MissingDriverError
9+
810
logger = logging.getLogger(__name__)
911

1012

13+
def _format_missing_driver_message(class_path: str) -> str:
14+
"""Format error message depending on whether the class path is a Jumpstarter driver."""
15+
# Extract package name from class path (first component)
16+
package_name = class_path.split(".")[0]
17+
18+
if class_path.startswith("jumpstarter_driver_"):
19+
return (
20+
f"Driver '{class_path}' is not installed.\n\n"
21+
"This usually indicates a version mismatch between your client and the exporter.\n"
22+
"Please try to update your client to the latest version and ensure the exporter "
23+
"has the correct version installed.\n"
24+
)
25+
else:
26+
return (
27+
f"Driver '{class_path}' is not installed.\n\n"
28+
"Please install the missing module:\n"
29+
f" pip install {package_name}\n\n"
30+
"or if using uv:\n"
31+
f" uv pip install {package_name}"
32+
)
33+
34+
1135
def cached_import(module_path, class_name):
1236
# Check whether module is loaded and fully initialized.
1337
if not (
@@ -40,36 +64,13 @@ def import_class(class_path: str, allow: list[str], unsafe: bool):
4064
try:
4165
return cached_import(module_path, class_name)
4266
except ModuleNotFoundError as e:
43-
module_name = str(e).split("'")[1] if "'" in str(e) else str(e).split()[-1]
44-
45-
is_jumpstarter_driver = unsafe or any(fnmatchcase(class_path, pattern) for pattern in allow)
46-
47-
if is_jumpstarter_driver:
48-
logger.error(
49-
"Missing Jumpstarter driver module '%s' for class '%s'. "
50-
"This usually indicates a version mismatch between your client and the exporter.",
51-
module_name,
52-
class_path,
53-
)
54-
raise ConnectionError(
55-
f"Missing Jumpstarter driver module '{module_name}'.\n\n"
56-
"This usually indicates a version mismatch between your client and the exporter.\n"
57-
"Please try to update your client to the latest version and ensure the exporter "
58-
"has the correct version installed.\n"
59-
) from e
60-
else:
61-
logger.error(
62-
"Missing Python module '%s' while importing '%s'. "
63-
"This module needs to be installed in your environment.",
64-
module_name,
65-
class_path,
66-
)
67-
raise ConnectionError(
68-
f"Missing Python module '{module_name}'.\n\n"
69-
"Please install the missing module:\n"
70-
f" pip install {module_name}\n\n"
71-
"or if using uv:\n"
72-
f" uv pip install {module_name}"
73-
) from e
67+
raise MissingDriverError(
68+
message=_format_missing_driver_message(class_path),
69+
class_path=class_path,
70+
) from e
7471
except AttributeError as e:
75-
raise ImportError(f"{module_path} doesn't have specified class {class_name}") from e
72+
# Module exists but class doesn't - treat in a similar way to missing module
73+
raise MissingDriverError(
74+
message=_format_missing_driver_message(class_path),
75+
class_path=class_path,
76+
) from e

python/packages/jumpstarter/jumpstarter/common/importlib_test.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
import pytest
22

3+
from .exceptions import MissingDriverError
34
from .importlib import import_class
45

56

67
def test_import_class():
78
import_class("os.open", [], True)
89

9-
with pytest.raises(ImportError):
10+
with pytest.raises(MissingDriverError):
1011
import_class("os.invalid", [], True)
1112

1213
with pytest.raises(ImportError):

python/packages/jumpstarter/jumpstarter/common/utils.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ def launch_shell(
108108
common_env = os.environ | {
109109
JUMPSTARTER_HOST: host,
110110
JMP_DRIVERS_ALLOW: "UNSAFE" if unsafe else ",".join(allow),
111+
"_JMP_SUPPRESS_DRIVER_WARNINGS": "1", # Already warned during client initialization
111112
}
112113

113114
if command:

python/packages/jumpstarter/jumpstarter/config/exporter.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from .common import ObjectMeta
1313
from .grpc import call_credentials
1414
from .tls import TLSConfigV1Alpha1
15-
from jumpstarter.common.exceptions import ConfigurationError
15+
from jumpstarter.common.exceptions import ConfigurationError, MissingDriverError
1616
from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials
1717
from jumpstarter.common.importlib import import_class
1818
from jumpstarter.driver import Driver
@@ -44,7 +44,12 @@ class ExporterConfigV1Alpha1DriverInstance(RootModel):
4444
def instantiate(self) -> Driver:
4545
match self.root:
4646
case ExporterConfigV1Alpha1DriverInstanceBase():
47-
driver_class = import_class(self.root.type, [], True)
47+
try:
48+
driver_class = import_class(self.root.type, [], True)
49+
except MissingDriverError:
50+
raise ConfigurationError(
51+
f"Driver '{self.root.type}' is not installed. Please check exporter configuration."
52+
) from None
4853

4954
children = {name: child.instantiate() for name, child in self.root.children.items()}
5055

0 commit comments

Comments
 (0)