Fix imports in test_embedded_ipython_code_executor.py

ag2ai · Jan 22, 2025 · f7a538c · f7a538c
1 parent 32eda33
commit f7a538c
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/test/coding/test_embedded_ipython_code_executor.py b/test/coding/test_embedded_ipython_code_executor.py
@@ -20,12 +20,6 @@
 )
 from autogen.coding.base import CodeBlock, CodeExecutor
 from autogen.coding.factory import CodeExecutorFactory
-from autogen.coding.jupyter import (
-    DockerJupyterServer,
-    EmbeddedIPythonCodeExecutor,
-    JupyterCodeExecutor,
-    LocalJupyterServer,
-)
 from autogen.import_utils import optional_import_block
 
 from ..conftest import MOCK_OPEN_AI_API_KEY
@@ -37,7 +31,18 @@
 
 classes_to_test = []
 with optional_import_block() as result:
+    import websocket  # noqa: F401
     from jupyter_client import KernelManager  # noqa: F401
+    from jupyter_client.kernelspec import KernelSpecManager  # noqa: F401
+
+    # This has to be here
+    # because autogen.coding.jupyter.helpers is written in a way that it will raise an ImportError
+    from autogen.coding.jupyter import (
+        DockerJupyterServer,
+        EmbeddedIPythonCodeExecutor,
+        JupyterCodeExecutor,
+        LocalJupyterServer,
+    )
 
 if result.is_successful: