From e65d81716223061bce0f1303248b19ac2353cbbf Mon Sep 17 00:00:00 2001
From: Kumaran Rajendhiran <kumaran@airt.ai>
Date: Wed, 22 Jan 2025 15:00:22 +0530
Subject: [PATCH] Only use third party modules in optional imports not autogen
 imports in test files (#602)

* Only use third party modules in optional imports not autogen imports in test/oai files

* Only use third party modules in optional imports not autogen imports in test/coding files

* Only use third party modules in optional imports not autogen imports in test/cache files

* Fix imports in test_embedded_ipython_code_executor.py

* Fix imports in leftover file

* Only use third party modules in optional imports not autogen imports in test/agentchat/contrib files

* Only use third party modules in optional imports not autogen imports in test/agentchat/contrib/vectordb files

* Only use third party modules in optional imports not autogen imports in test/agentchat/contrib/retrievechat files

* Only use third party modules in optional imports not autogen imports in test/agentchat/contrib/graph_rag files

* Only use third party modules in optional imports not autogen imports in test/agentchat/contrib/capabilities files

* tests fixed

---------

Co-authored-by: Davor Runje <davor@airt.ai>
---
 .../graph_rag/falkor_graph_query_engine.py    |   4 +-
 .../contrib/graph_rag/graph_query_engine.py   |   2 +-
 .../graph_rag/neo4j_graph_query_engine.py     |   2 +-
 .../agentchat/contrib/vectordb/pgvectordb.py  |   6 +-
 autogen/coding/jupyter/__init__.py            |   7 -
 .../coding/jupyter/docker_jupyter_server.py   |   2 +
 .../jupyter/embedded_ipython_code_executor.py |   2 +
 autogen/coding/jupyter/helpers.py             |  27 --
 autogen/coding/jupyter/import_utils.py        |  78 ++++
 .../coding/jupyter/local_jupyter_server.py    |   2 +
 autogen/formatting_utils.py                   |  12 +-
 autogen/import_utils.py                       |  76 ++--
 autogen/runtime_logging.py                    |   6 +-
 pyproject.toml                                |   3 +
 .../capabilities/test_teachable_agent.py      |   4 +-
 .../contrib/capabilities/test_transforms.py   |   4 +-
 .../capabilities/test_vision_capability.py    |   2 +-
 .../graph_rag/test_native_neo4j_graph_rag.py  |  12 +-
 .../contrib/graph_rag/test_neo4j_graph_rag.py |  11 +-
 .../test_pgvector_retrievechat.py             |   8 +-
 .../retrievechat/test_qdrant_retrievechat.py  |  10 +-
 test/agentchat/contrib/test_img_utils.py      |  22 +-
 test/agentchat/contrib/test_llava.py          |   3 +-
 test/agentchat/contrib/test_lmm.py            |   6 +-
 test/agentchat/contrib/test_web_surfer.py     |   8 +-
 .../contrib/vectordb/test_chromadb.py         |   2 +-
 .../contrib/vectordb/test_mongodb.py          |   2 +-
 .../contrib/vectordb/test_pgvectordb.py       |   2 +-
 .../agentchat/contrib/vectordb/test_qdrant.py |   6 +-
 test/cache/test_cosmos_db_cache.py            |   3 +-
 .../test_embedded_ipython_code_executor.py    | 398 +++++++++---------
 test/oai/test_anthropic.py                    |   2 +-
 test/oai/test_cohere.py                       |   3 +-
 test/oai/test_gemini.py                       |   6 +-
 test/oai/test_groq.py                         |   3 +-
 test/oai/test_mistral.py                      |   2 +-
 test/test_import.py                           |  84 ++++
 test/test_import_utils.py                     |  36 +-
 38 files changed, 507 insertions(+), 361 deletions(-)
 delete mode 100644 autogen/coding/jupyter/helpers.py
 create mode 100644 autogen/coding/jupyter/import_utils.py
 create mode 100644 test/test_import.py

diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
index d9ecb9330..79bcb5b5c 100644
--- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
+++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
@@ -28,8 +28,8 @@ def __init__(
         name: str,
         host: str = "127.0.0.1",
         port: int = 6379,
-        username: str | None = None,
-        password: str | None = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
         model: Optional["GenerativeModel"] = None,
         ontology: Optional["Ontology"] = None,
     ):
diff --git a/autogen/agentchat/contrib/graph_rag/graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/graph_query_engine.py
index b474902a6..6406cb9aa 100644
--- a/autogen/agentchat/contrib/graph_rag/graph_query_engine.py
+++ b/autogen/agentchat/contrib/graph_rag/graph_query_engine.py
@@ -28,7 +28,7 @@ class GraphQueryEngine(Protocol):
     This interface defines the basic methods for graph-based RAG.
     """
 
-    def init_db(self, input_doc: list[Document] | None = None):
+    def init_db(self, input_doc: Optional[list[Document]] = None):
         """This method initializes graph database with the input documents or records.
         Usually, it takes the following steps,
         1. connecting to a graph database.
diff --git a/autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py
index b86dbf287..cc9afff17 100644
--- a/autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py
+++ b/autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py
@@ -94,7 +94,7 @@ def __init__(
         self.schema = schema
         self.strict = strict
 
-    def init_db(self, input_doc: list[Document] | None = None):
+    def init_db(self, input_doc: Optional[list[Document]] = None):
         """Build the knowledge graph with input documents."""
         self.documents = self._load_doc(input_doc)
 
diff --git a/autogen/agentchat/contrib/vectordb/pgvectordb.py b/autogen/agentchat/contrib/vectordb/pgvectordb.py
index 0b87b38e6..fdcd1e918 100644
--- a/autogen/agentchat/contrib/vectordb/pgvectordb.py
+++ b/autogen/agentchat/contrib/vectordb/pgvectordb.py
@@ -10,7 +10,6 @@
 from typing import Callable, Optional, Union
 
 import numpy as np
-from sentence_transformers import SentenceTransformer
 
 from ....import_utils import optional_import_block, require_optional_import
 from .base import Document, ItemID, QueryResults, VectorDB
@@ -20,12 +19,13 @@
     import pgvector  # noqa: F401
     import psycopg
     from pgvector.psycopg import register_vector
+    from sentence_transformers import SentenceTransformer
 
 PGVECTOR_MAX_BATCH_SIZE = os.environ.get("PGVECTOR_MAX_BATCH_SIZE", 40000)
 logger = get_logger(__name__)
 
 
-@require_optional_import("psycopg", "retrievechat-pgvector")
+@require_optional_import(["psycopg", "sentence_transformers"], "retrievechat-pgvector")
 class Collection:
     """A Collection object for PGVector.
 
@@ -539,7 +539,7 @@ def create_collection(
         cursor.close()
 
 
-@require_optional_import(["pgvector", "psycopg"], "retrievechat-pgvector")
+@require_optional_import(["pgvector", "psycopg", "sentence_transformers"], "retrievechat-pgvector")
 class PGVectorDB(VectorDB):
     """A vector database that uses PGVector as the backend."""
 
diff --git a/autogen/coding/jupyter/__init__.py b/autogen/coding/jupyter/__init__.py
index a15fe09ff..729fd08e8 100644
--- a/autogen/coding/jupyter/__init__.py
+++ b/autogen/coding/jupyter/__init__.py
@@ -4,13 +4,6 @@
 #
 # Original portions of this file are derived from https://github.com/microsoft/autogen under the MIT License.
 # SPDX-License-Identifier: MIT
-from .helpers import is_jupyter_kernel_gateway_installed
-
-if not is_jupyter_kernel_gateway_installed():
-    raise ImportError(
-        "jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
-    )
-
 
 from .base import JupyterConnectable, JupyterConnectionInfo
 from .docker_jupyter_server import DockerJupyterServer
diff --git a/autogen/coding/jupyter/docker_jupyter_server.py b/autogen/coding/jupyter/docker_jupyter_server.py
index 2191a1b22..f3307eb68 100644
--- a/autogen/coding/jupyter/docker_jupyter_server.py
+++ b/autogen/coding/jupyter/docker_jupyter_server.py
@@ -24,9 +24,11 @@
 
 from ..docker_commandline_code_executor import _wait_for_ready
 from .base import JupyterConnectable, JupyterConnectionInfo
+from .import_utils import require_jupyter_kernel_gateway_installed
 from .jupyter_client import JupyterClient
 
 
+@require_jupyter_kernel_gateway_installed()
 class DockerJupyterServer(JupyterConnectable):
     DEFAULT_DOCKERFILE = """FROM quay.io/jupyter/docker-stacks-foundation
 
diff --git a/autogen/coding/jupyter/embedded_ipython_code_executor.py b/autogen/coding/jupyter/embedded_ipython_code_executor.py
index b02c3af78..6dc30b859 100644
--- a/autogen/coding/jupyter/embedded_ipython_code_executor.py
+++ b/autogen/coding/jupyter/embedded_ipython_code_executor.py
@@ -18,6 +18,7 @@
 from ...import_utils import optional_import_block, require_optional_import
 from ..base import CodeBlock, CodeExtractor, IPythonCodeResult
 from ..markdown_code_extractor import MarkdownCodeExtractor
+from .import_utils import require_jupyter_kernel_gateway_installed
 
 with optional_import_block():
     from jupyter_client import KernelManager  # type: ignore[attr-defined]
@@ -27,6 +28,7 @@
 
 
 @require_optional_import("jupyter_client", "jupyter-executor")
+@require_jupyter_kernel_gateway_installed()
 class EmbeddedIPythonCodeExecutor(BaseModel):
     """(Experimental) A code executor class that executes code statefully using an embedded
     IPython kernel managed by this class.
diff --git a/autogen/coding/jupyter/helpers.py b/autogen/coding/jupyter/helpers.py
deleted file mode 100644
index 5f8cf1cc6..000000000
--- a/autogen/coding/jupyter/helpers.py
+++ /dev/null
@@ -1,27 +0,0 @@
-# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
-#
-# SPDX-License-Identifier: Apache-2.0
-
-import subprocess
-from logging import getLogger
-
-logger = getLogger(__name__)
-
-__all__ = ["is_jupyter_kernel_gateway_installed"]
-
-
-def is_jupyter_kernel_gateway_installed() -> bool:
-    """Check if jupyter-kernel-gateway is installed."""
-    try:
-        subprocess.run(
-            ["jupyter", "kernelgateway", "--version"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            check=True,
-        )
-        return True
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        logger.warning(
-            "jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
-        )
-        return False
diff --git a/autogen/coding/jupyter/import_utils.py b/autogen/coding/jupyter/import_utils.py
new file mode 100644
index 000000000..fe0c83526
--- /dev/null
+++ b/autogen/coding/jupyter/import_utils.py
@@ -0,0 +1,78 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import subprocess
+from functools import lru_cache
+from logging import getLogger
+from typing import Callable, TypeVar
+
+from ...import_utils import patch_object
+
+logger = getLogger(__name__)
+
+__all__ = ["require_jupyter_kernel_gateway_installed", "skip_on_missing_jupyter_kernel_gateway"]
+
+
+@lru_cache()
+def is_jupyter_kernel_gateway_installed() -> bool:
+    """Check if jupyter-kernel-gateway is installed."""
+    try:
+        subprocess.run(
+            ["jupyter", "kernelgateway", "--version"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            check=True,
+        )
+        return True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        logger.warning(
+            "jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
+        )
+        return False
+
+
+T = TypeVar("T")
+
+
+def require_jupyter_kernel_gateway_installed() -> Callable[[T], T]:
+    """Decorator to handle optional module dependencies
+
+    Args:
+        modules: Module name or list of module names required
+        dep_target: Target name for pip installation (e.g. 'test' in pip install ag2[test])
+    """
+    if is_jupyter_kernel_gateway_installed():
+
+        def decorator(o: T) -> T:
+            return o
+    else:
+
+        def decorator(o: T) -> T:
+            return patch_object(o, missing_modules=[], dep_target="jupyter-executor")
+
+    return decorator
+
+
+def skip_on_missing_jupyter_kernel_gateway() -> Callable[[T], T]:
+    """Decorator to skip a test if an optional module is missing
+
+    Args:
+        module: Module name
+        dep_target: Target name for pip installation (e.g. 'test' in pip install ag2[test])
+    """
+
+    if is_jupyter_kernel_gateway_installed():
+
+        def decorator(o: T) -> T:
+            return o
+    else:
+
+        def decorator(o: T) -> T:
+            import pytest
+
+            return pytest.mark.skip(
+                reason="jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
+            )(o)  # type: ignore[return-value]
+
+    return decorator
diff --git a/autogen/coding/jupyter/local_jupyter_server.py b/autogen/coding/jupyter/local_jupyter_server.py
index 518b2171f..0e77bdccb 100644
--- a/autogen/coding/jupyter/local_jupyter_server.py
+++ b/autogen/coding/jupyter/local_jupyter_server.py
@@ -20,9 +20,11 @@
     from typing_extensions import Self
 
 from .base import JupyterConnectable, JupyterConnectionInfo
+from .import_utils import require_jupyter_kernel_gateway_installed
 from .jupyter_client import JupyterClient
 
 
+@require_jupyter_kernel_gateway_installed()
 class LocalJupyterServer(JupyterConnectable):
     class GenerateToken:
         pass
diff --git a/autogen/formatting_utils.py b/autogen/formatting_utils.py
index eee090a24..5d10ae74b 100644
--- a/autogen/formatting_utils.py
+++ b/autogen/formatting_utils.py
@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Literal
+from typing import Literal, Optional
 
 from .import_utils import optional_import_block
 
@@ -70,12 +70,12 @@
 
     def colored(
         text: object,
-        color: Color | None = None,
-        on_color: Highlight | None = None,
-        attrs: Iterable[Attribute] | None = None,
+        color: Optional[Color] = None,
+        on_color: Optional[Highlight] = None,
+        attrs: Optional[Iterable[Attribute]] = None,
         *,
-        no_color: bool | None = None,
-        force_color: bool | None = None,
+        no_color: Optional[bool] = None,
+        force_color: Optional[bool] = None,
     ) -> str:
         return str(text)
 
diff --git a/autogen/import_utils.py b/autogen/import_utils.py
index 8a526f64b..1f8e408a2 100644
--- a/autogen/import_utils.py
+++ b/autogen/import_utils.py
@@ -8,19 +8,19 @@
 from contextlib import contextmanager
 from functools import wraps
 from logging import getLogger
-from typing import Any, Callable, Generator, Iterable, Optional, Type, TypeVar, Union
+from typing import Any, Callable, Generator, Generic, Iterable, Optional, Type, TypeVar, Union
 
-__all__ = ["optional_import_block", "require_optional_import", "skip_on_missing_imports"]
+__all__ = ["optional_import_block", "patch_object", "require_optional_import", "skip_on_missing_imports"]
 
 logger = getLogger(__name__)
 
 
 class Result:
-    def __init__(self):
+    def __init__(self) -> None:
         self._failed: Optional[bool] = None
 
     @property
-    def is_successful(self):
+    def is_successful(self) -> bool:
         if self._failed is None:
             raise ValueError("Result not set")
         return not self._failed
@@ -50,7 +50,7 @@ def optional_import_block() -> Generator[Result, None, None]:
         result._failed = True
 
 
-def get_missing_imports(modules: Union[str, Iterable[str]]):
+def get_missing_imports(modules: Union[str, Iterable[str]]) -> list[str]:
     """Get missing modules from a list of module names
 
     Args:
@@ -69,13 +69,13 @@ def get_missing_imports(modules: Union[str, Iterable[str]]):
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-class PatchObject(ABC):
+class PatchObject(ABC, Generic[T]):
     def __init__(self, o: T, missing_modules: Iterable[str], dep_target: str):
         if not self.accept(o):
             raise ValueError(f"Cannot patch object of type {type(o)}")
 
         self.o = o
-        self.missing_modules = missing_modules
+        self.missing_modules = list(missing_modules)
         self.dep_target = dep_target
 
     @classmethod
@@ -107,22 +107,22 @@ def copy_metadata(self, retval: T) -> None:
         if hasattr(o, "__doc__"):
             retval.__doc__ = o.__doc__
         if hasattr(o, "__name__"):
-            retval.__name__ = o.__name__
+            retval.__name__ = o.__name__  # type: ignore[attr-defined]
         if hasattr(o, "__module__"):
             retval.__module__ = o.__module__
 
-    _registry: list[Type["PatchObject"]] = []
+    _registry: list[Type["PatchObject[Any]"]] = []
 
     @classmethod
-    def register(cls) -> Callable[[Type["PatchObject"]], Type["PatchObject"]]:
-        def decorator(subclass: Type["PatchObject"]):
+    def register(cls) -> Callable[[Type["PatchObject[Any]"]], Type["PatchObject[Any]"]]:
+        def decorator(subclass: Type["PatchObject[Any]"]) -> Type["PatchObject[Any]"]:
             cls._registry.append(subclass)
             return subclass
 
         return decorator
 
     @classmethod
-    def create(cls, o: T, *, missing_modules: Iterable[str], dep_target: str) -> Optional["PatchObject"]:
+    def create(cls, o: T, *, missing_modules: Iterable[str], dep_target: str) -> Optional["PatchObject[T]"]:
         # print(f"{cls._registry=}")
         for subclass in cls._registry:
             if subclass.accept(o):
@@ -131,7 +131,7 @@ def create(cls, o: T, *, missing_modules: Iterable[str], dep_target: str) -> Opt
 
 
 @PatchObject.register()
-class PatchCallable(PatchObject):
+class PatchCallable(PatchObject[F]):
     @classmethod
     def accept(cls, o: Any) -> bool:
         return inspect.isfunction(o) or inspect.ismethod(o)
@@ -139,39 +139,39 @@ def accept(cls, o: Any) -> bool:
     def patch(self) -> F:
         f: Callable[..., Any] = self.o
 
-        @wraps(f.__call__)
-        def _call(*args, **kwargs):
+        @wraps(f.__call__)  # type: ignore[operator]
+        def _call(*args: Any, **kwargs: Any) -> Any:
             raise ImportError(self.msg)
 
-        self.copy_metadata(_call)
+        self.copy_metadata(_call)  # type: ignore[arg-type]
 
-        return _call
+        return _call  # type: ignore[return-value]
 
 
 @PatchObject.register()
-class PatchStatic(PatchObject):
+class PatchStatic(PatchObject[F]):
     @classmethod
     def accept(cls, o: Any) -> bool:
         # return inspect.ismethoddescriptor(o)
         return isinstance(o, staticmethod)
 
     def patch(self) -> F:
-        f: Callable[..., Any] = self.o.__func__
+        f: Callable[..., Any] = self.o.__func__  # type: ignore[attr-defined]
 
         @wraps(f)
-        def _call(*args, **kwargs):
+        def _call(*args: Any, **kwargs: Any) -> Any:
             raise ImportError(self.msg)
 
-        self.copy_metadata(_call)
+        self.copy_metadata(_call)  # type: ignore[arg-type]
 
-        return staticmethod(_call)
+        return staticmethod(_call)  # type: ignore[return-value]
 
     def get_object_with_metadata(self) -> Any:
-        return self.o.__func__
+        return self.o.__func__  # type: ignore[attr-defined]
 
 
 @PatchObject.register()
-class PatchInit(PatchObject):
+class PatchInit(PatchObject[F]):
     @classmethod
     def accept(cls, o: Any) -> bool:
         return inspect.ismethoddescriptor(o) and o.__name__ == "__init__"
@@ -180,19 +180,19 @@ def patch(self) -> F:
         f: Callable[..., Any] = self.o
 
         @wraps(f)
-        def _call(*args, **kwargs):
+        def _call(*args: Any, **kwargs: Any) -> Any:
             raise ImportError(self.msg)
 
-        self.copy_metadata(_call)
+        self.copy_metadata(_call)  # type: ignore[arg-type]
 
-        return staticmethod(_call)
+        return staticmethod(_call)  # type: ignore[return-value]
 
     def get_object_with_metadata(self) -> Any:
         return self.o
 
 
 @PatchObject.register()
-class PatchProperty(PatchObject):
+class PatchProperty(PatchObject[Any]):
     @classmethod
     def accept(cls, o: Any) -> bool:
         return inspect.isdatadescriptor(o) and hasattr(o, "fget")
@@ -203,7 +203,7 @@ def patch(self) -> property:
         f: Callable[..., Any] = self.o.fget
 
         @wraps(f)
-        def _call(*args, **kwargs):
+        def _call(*args: Any, **kwargs: Any) -> Any:
             raise ImportError(self.msg)
 
         self.copy_metadata(_call)
@@ -215,12 +215,12 @@ def get_object_with_metadata(self) -> Any:
 
 
 @PatchObject.register()
-class PatchClass(PatchObject):
+class PatchClass(PatchObject[Type[Any]]):
     @classmethod
     def accept(cls, o: Any) -> bool:
         return inspect.isclass(o)
 
-    def patch(self) -> F:
+    def patch(self) -> Type[Any]:
         # Patch __init__ method if possible
 
         for name, member in inspect.getmembers(self.o):
@@ -229,7 +229,11 @@ def patch(self) -> F:
             patched = patch_object(
                 member, missing_modules=self.missing_modules, dep_target=self.dep_target, fail_if_not_patchable=False
             )
-            setattr(self.o, name, patched)
+            print(f"Patching {name=}, {member=}, {patched=}")
+            try:
+                setattr(self.o, name, patched)
+            except AttributeError:
+                pass
 
         return self.o
 
@@ -283,12 +287,8 @@ def decorator(o: T) -> T:
         def decorator(o: T) -> T:
             import pytest
 
-            @pytest.mark.skip(
+            return pytest.mark.skip(  # type: ignore[return-value]
                 f"Missing module{'s' if len(missing_modules) > 1 else ''}: {', '.join(missing_modules)}. Install using 'pip install ag2[{dep_target}]'"
-            )
-            def _skip(*args, **kwargs):
-                pass
-
-            return _skip
+            )(o)
 
     return decorator
diff --git a/autogen/runtime_logging.py b/autogen/runtime_logging.py
index af2f7021f..f87e33526 100644
--- a/autogen/runtime_logging.py
+++ b/autogen/runtime_logging.py
@@ -9,7 +9,7 @@
 import logging
 import sqlite3
 import uuid
-from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, TypeVar
 
 from openai import AzureOpenAI, OpenAI
 from openai.types.chat import ChatCompletion
@@ -38,9 +38,9 @@
 
 
 def start(
-    logger: BaseLogger | None = None,
+    logger: Optional[BaseLogger] = None,
     logger_type: Literal["sqlite", "file"] = "sqlite",
-    config: dict[str, Any] | None = None,
+    config: Optional[dict[str, Any]] = None,
 ) -> str:
     """Start logging for the runtime.
 
diff --git a/pyproject.toml b/pyproject.toml
index 1d5a19c9e..806d1757f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -345,6 +345,7 @@ files = [
     "autogen/interop",
     "autogen/agentchat/realtime_agent",
     "autogen/messages",
+    "autogen/import_utils.py",
     "test/test_pydantic.py",
     "test/io",
     "test/tools",
@@ -352,6 +353,8 @@ files = [
     "test/agentchat/realtime_agent",
     "test/messages",
     "test/conftest.py",
+    "test/test_import_utils.py",
+    "test/test_import.py",
 ]
 
 exclude = [
diff --git a/test/agentchat/contrib/capabilities/test_teachable_agent.py b/test/agentchat/contrib/capabilities/test_teachable_agent.py
index dd626ba4c..59d9a7ab6 100755
--- a/test/agentchat/contrib/capabilities/test_teachable_agent.py
+++ b/test/agentchat/contrib/capabilities/test_teachable_agent.py
@@ -9,13 +9,15 @@
 import pytest
 
 from autogen import ConversableAgent
+from autogen.agentchat.contrib.capabilities.teachability import Teachability
 from autogen.formatting_utils import colored
 from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.capabilities.teachability import Teachability
+    import chromadb  # noqa: F401
+
 
 skip = not result.is_successful
 
diff --git a/test/agentchat/contrib/capabilities/test_transforms.py b/test/agentchat/contrib/capabilities/test_transforms.py
index 24f9ca8fc..a45a6d116 100644
--- a/test/agentchat/contrib/capabilities/test_transforms.py
+++ b/test/agentchat/contrib/capabilities/test_transforms.py
@@ -106,9 +106,11 @@ def get_messages_with_names_post_filtered() -> list[dict]:
 def get_text_compressors() -> list[TextCompressor]:
     compressors: list[TextCompressor] = [_MockTextCompressor()]
     with optional_import_block() as result:
-        from autogen.agentchat.contrib.capabilities.text_compressors import LLMLingua
+        import llmlingua  # noqa: F401
 
     if result.is_successful:
+        from autogen.agentchat.contrib.capabilities.text_compressors import LLMLingua
+
         compressors.append(LLMLingua())
 
     return compressors
diff --git a/test/agentchat/contrib/capabilities/test_vision_capability.py b/test/agentchat/contrib/capabilities/test_vision_capability.py
index 3006b6786..dad08ab0f 100644
--- a/test/agentchat/contrib/capabilities/test_vision_capability.py
+++ b/test/agentchat/contrib/capabilities/test_vision_capability.py
@@ -9,13 +9,13 @@
 
 import pytest
 
+from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
 from autogen.agentchat.conversable_agent import ConversableAgent
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
     from PIL import Image  # noqa: F401
 
-    from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
 
 skip_test = not result.is_successful
 
diff --git a/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py b/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py
index 12932c551..a303578a2 100644
--- a/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py
+++ b/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py
@@ -7,16 +7,18 @@
 
 import pytest
 
+from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
+from autogen.agentchat.contrib.graph_rag.neo4j_native_graph_query_engine import (
+    GraphStoreQueryResult,
+    Neo4jNativeGraphQueryEngine,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import reason
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
-    from autogen.agentchat.contrib.graph_rag.neo4j_native_graph_query_engine import (
-        GraphStoreQueryResult,
-        Neo4jNativeGraphQueryEngine,
-    )
+    from neo4j import GraphDatabase  # noqa: F401
+    from neo4j_graphrag.embeddings import Embedder  # noqa: F401
 
 
 skip = not result.is_successful
diff --git a/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py b/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py
index 22e8d49d5..79c6faae0 100644
--- a/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py
+++ b/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py
@@ -10,16 +10,17 @@
 
 import pytest
 
+from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
+from autogen.agentchat.contrib.graph_rag.neo4j_graph_query_engine import (
+    GraphStoreQueryResult,
+    Neo4jGraphQueryEngine,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import reason
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
-    from autogen.agentchat.contrib.graph_rag.neo4j_graph_query_engine import (
-        GraphStoreQueryResult,
-        Neo4jGraphQueryEngine,
-    )
+    from llama_index.core import PropertyGraphIndex  # noqa: F401
 
 
 skip = not result.is_successful
diff --git a/test/agentchat/contrib/retrievechat/test_pgvector_retrievechat.py b/test/agentchat/contrib/retrievechat/test_pgvector_retrievechat.py
index c79d9020d..dadf394ab 100644
--- a/test/agentchat/contrib/retrievechat/test_pgvector_retrievechat.py
+++ b/test/agentchat/contrib/retrievechat/test_pgvector_retrievechat.py
@@ -12,16 +12,18 @@
 from sentence_transformers import SentenceTransformer
 
 from autogen import AssistantAgent
+from autogen.agentchat.contrib.retrieve_user_proxy_agent import (
+    RetrieveUserProxyAgent,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
 
 with optional_import_block() as result:
+    import chromadb  # noqa: F401
     import pgvector  # noqa: F401
+    from IPython import get_ipython  # noqa: F401
 
-    from autogen.agentchat.contrib.retrieve_user_proxy_agent import (
-        RetrieveUserProxyAgent,
-    )
 
 skip = not result.is_successful
 
diff --git a/test/agentchat/contrib/retrievechat/test_qdrant_retrievechat.py b/test/agentchat/contrib/retrievechat/test_qdrant_retrievechat.py
index 696b20b80..67af46885 100755
--- a/test/agentchat/contrib/retrievechat/test_qdrant_retrievechat.py
+++ b/test/agentchat/contrib/retrievechat/test_qdrant_retrievechat.py
@@ -12,6 +12,11 @@
 import pytest
 
 from autogen import AssistantAgent
+from autogen.agentchat.contrib.qdrant_retrieve_user_proxy_agent import (
+    QdrantRetrieveUserProxyAgent,
+    create_qdrant_from_dir,
+    query_qdrant,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
@@ -20,11 +25,6 @@
     import fastembed  # noqa: F401
     from qdrant_client import QdrantClient
 
-    from autogen.agentchat.contrib.qdrant_retrieve_user_proxy_agent import (
-        QdrantRetrieveUserProxyAgent,
-        create_qdrant_from_dir,
-        query_qdrant,
-    )
 
 QDRANT_INSTALLED = result.is_successful
 
diff --git a/test/agentchat/contrib/test_img_utils.py b/test/agentchat/contrib/test_img_utils.py
index b6fcf9b14..98add69b5 100755
--- a/test/agentchat/contrib/test_img_utils.py
+++ b/test/agentchat/contrib/test_img_utils.py
@@ -11,25 +11,25 @@
 import unittest
 from unittest.mock import patch
 
+import numpy as np
 import pytest
 import requests
 
+from autogen.agentchat.contrib.img_utils import (
+    convert_base64_to_data_uri,
+    extract_img_paths,
+    get_image_data,
+    get_pil_image,
+    gpt4v_formatter,
+    llava_formatter,
+    message_formatter_pil_to_b64,
+    num_tokens_from_gpt_image,
+)
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
-    import numpy as np
     from PIL import Image
 
-    from autogen.agentchat.contrib.img_utils import (
-        convert_base64_to_data_uri,
-        extract_img_paths,
-        get_image_data,
-        get_pil_image,
-        gpt4v_formatter,
-        llava_formatter,
-        message_formatter_pil_to_b64,
-        num_tokens_from_gpt_image,
-    )
 
 skip = not result.is_successful
 
diff --git a/test/agentchat/contrib/test_llava.py b/test/agentchat/contrib/test_llava.py
index 38041203a..37ae0b299 100755
--- a/test/agentchat/contrib/test_llava.py
+++ b/test/agentchat/contrib/test_llava.py
@@ -11,12 +11,13 @@
 
 import pytest
 
+from autogen.agentchat.contrib.llava_agent import LLaVAAgent, _llava_call_binary_with_config, llava_call
 from autogen.import_utils import optional_import_block
 
 from ...conftest import MOCK_OPEN_AI_API_KEY
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.llava_agent import LLaVAAgent, _llava_call_binary_with_config, llava_call
+    import replicate  # noqa: F401
 
 skip = not result.is_successful
 
diff --git a/test/agentchat/contrib/test_lmm.py b/test/agentchat/contrib/test_lmm.py
index a65962d3d..7a87bd469 100755
--- a/test/agentchat/contrib/test_lmm.py
+++ b/test/agentchat/contrib/test_lmm.py
@@ -12,14 +12,16 @@
 import pytest
 
 import autogen
+from autogen.agentchat.contrib.img_utils import get_pil_image
+from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent
 from autogen.agentchat.conversable_agent import ConversableAgent
 from autogen.import_utils import optional_import_block
 
 from ...conftest import MOCK_OPEN_AI_API_KEY
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.img_utils import get_pil_image
-    from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent
+    from PIL import Image  # noqa: F401
+
 
 skip = not result.is_successful
 
diff --git a/test/agentchat/contrib/test_web_surfer.py b/test/agentchat/contrib/test_web_surfer.py
index edd9c38d0..65f467817 100755
--- a/test/agentchat/contrib/test_web_surfer.py
+++ b/test/agentchat/contrib/test_web_surfer.py
@@ -12,6 +12,7 @@
 import pytest
 
 from autogen import UserProxyAgent
+from autogen.agentchat.contrib.web_surfer import WebSurferAgent
 from autogen.import_utils import optional_import_block
 
 from ...conftest import MOCK_OPEN_AI_API_KEY, Credentials
@@ -21,7 +22,12 @@
 BING_QUERY = "Microsoft"
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.web_surfer import WebSurferAgent
+    import markdownify  # noqa: F401
+    import pathvalidate  # noqa: F401
+    import pdfminer  # noqa: F401
+    import requests  # noqa: F401
+    from bs4 import BeautifulSoup  # noqa: F401
+
 
 skip_all = not result.is_successful
 
diff --git a/test/agentchat/contrib/vectordb/test_chromadb.py b/test/agentchat/contrib/vectordb/test_chromadb.py
index 3e2e1ca0a..c45f7f117 100644
--- a/test/agentchat/contrib/vectordb/test_chromadb.py
+++ b/test/agentchat/contrib/vectordb/test_chromadb.py
@@ -9,6 +9,7 @@
 
 import pytest
 
+from autogen.agentchat.contrib.vectordb.chromadb import ChromaVectorDB
 from autogen.import_utils import optional_import_block
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
@@ -18,7 +19,6 @@
     import chromadb.errors
     import sentence_transformers  # noqa: F401
 
-    from autogen.agentchat.contrib.vectordb.chromadb import ChromaVectorDB
 
 skip = not result.is_successful
 
diff --git a/test/agentchat/contrib/vectordb/test_mongodb.py b/test/agentchat/contrib/vectordb/test_mongodb.py
index 66853fe8e..76877e610 100644
--- a/test/agentchat/contrib/vectordb/test_mongodb.py
+++ b/test/agentchat/contrib/vectordb/test_mongodb.py
@@ -12,13 +12,13 @@
 import pytest
 
 from autogen.agentchat.contrib.vectordb.base import Document
+from autogen.agentchat.contrib.vectordb.mongodb import MongoDBAtlasVectorDB
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
     import pymongo  # noqa: F401
     import sentence_transformers  # noqa: F401
 
-    from autogen.agentchat.contrib.vectordb.mongodb import MongoDBAtlasVectorDB
 
 if not result.is_successful:
     # To display warning in pyproject.toml [tool.pytest.ini_options] set log_cli = true
diff --git a/test/agentchat/contrib/vectordb/test_pgvectordb.py b/test/agentchat/contrib/vectordb/test_pgvectordb.py
index 9083ad954..47b6e6e50 100644
--- a/test/agentchat/contrib/vectordb/test_pgvectordb.py
+++ b/test/agentchat/contrib/vectordb/test_pgvectordb.py
@@ -10,6 +10,7 @@
 
 import pytest
 
+from autogen.agentchat.contrib.vectordb.pgvectordb import PGVectorDB
 from autogen.import_utils import optional_import_block
 
 from ....conftest import reason
@@ -19,7 +20,6 @@
     import psycopg
     import sentence_transformers  # noqa: F401
 
-    from autogen.agentchat.contrib.vectordb.pgvectordb import PGVectorDB
 
 skip = not result.is_successful
 
diff --git a/test/agentchat/contrib/vectordb/test_qdrant.py b/test/agentchat/contrib/vectordb/test_qdrant.py
index c7dd1385b..20870eddd 100644
--- a/test/agentchat/contrib/vectordb/test_qdrant.py
+++ b/test/agentchat/contrib/vectordb/test_qdrant.py
@@ -6,19 +6,19 @@
 # SPDX-License-Identifier: MIT
 import os
 import sys
+import uuid
 
 import pytest
 
+from autogen.agentchat.contrib.vectordb.qdrant import QdrantVectorDB
 from autogen.import_utils import optional_import_block
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 
 with optional_import_block() as result:
-    import uuid
-
+    from fastembed import TextEmbedding  # noqa: F401
     from qdrant_client import QdrantClient
 
-    from autogen.agentchat.contrib.vectordb.qdrant import QdrantVectorDB
 
 skip = not result.is_successful
 
diff --git a/test/cache/test_cosmos_db_cache.py b/test/cache/test_cosmos_db_cache.py
index 268237629..056cebc11 100644
--- a/test/cache/test_cosmos_db_cache.py
+++ b/test/cache/test_cosmos_db_cache.py
@@ -10,12 +10,13 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
+from autogen.cache.cosmos_db_cache import CosmosDBCache
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
+    from azure.cosmos import CosmosClient  # noqa: F401
     from azure.cosmos.exceptions import CosmosResourceNotFoundError
 
-    from autogen.cache.cosmos_db_cache import CosmosDBCache
 
 skip_test = not result.is_successful
 
diff --git a/test/coding/test_embedded_ipython_code_executor.py b/test/coding/test_embedded_ipython_code_executor.py
index dd8e0ff29..08602883a 100644
--- a/test/coding/test_embedded_ipython_code_executor.py
+++ b/test/coding/test_embedded_ipython_code_executor.py
@@ -20,228 +20,214 @@
 )
 from autogen.coding.base import CodeBlock, CodeExecutor
 from autogen.coding.factory import CodeExecutorFactory
-from autogen.import_utils import optional_import_block
+from autogen.coding.jupyter import (
+    DockerJupyterServer,
+    EmbeddedIPythonCodeExecutor,
+    JupyterCodeExecutor,
+    LocalJupyterServer,
+)
+from autogen.coding.jupyter.import_utils import skip_on_missing_jupyter_kernel_gateway
+from autogen.import_utils import optional_import_block, skip_on_missing_imports
 
 from ..conftest import MOCK_OPEN_AI_API_KEY
 
-if not is_docker_running() or not decide_use_docker(use_docker=None):
-    skip_docker_test = True
-else:
-    skip_docker_test = False
+# needed for skip_on_missing_imports to work
+with optional_import_block():
+    import ipykernel  # noqa: F401
+    import jupyter_client  # noqa: F401
+    import requests  # noqa: F401
+    import websocket  # noqa: F401
 
-classes_to_test = []
-with optional_import_block() as result:
-    from autogen.coding.jupyter import (
-        DockerJupyterServer,
-        EmbeddedIPythonCodeExecutor,
-        JupyterCodeExecutor,
-        LocalJupyterServer,
-    )
-
-if result.is_successful:
-
-    class DockerJupyterExecutor(JupyterCodeExecutor):
-        def __init__(self, **kwargs):
-            jupyter_server = DockerJupyterServer()
-            super().__init__(jupyter_server=jupyter_server, **kwargs)
-
-    class LocalJupyterCodeExecutor(JupyterCodeExecutor):
-        def __init__(self, **kwargs):
-            jupyter_server = LocalJupyterServer()
-            super().__init__(jupyter_server=jupyter_server, **kwargs)
-
-    # Skip on windows due to kernelgateway bug https://github.com/jupyter-server/kernel_gateway/issues/398
-    if sys.platform == "win32":
-        classes_to_test = [EmbeddedIPythonCodeExecutor]
-    else:
-        classes_to_test = [EmbeddedIPythonCodeExecutor, LocalJupyterCodeExecutor]
-
-    if not skip_docker_test:
-        classes_to_test.append(DockerJupyterExecutor)
-
-skip = not result.is_successful
-skip_reason = (
-    ""
-    if result.is_successful
-    else "Dependencies for EmbeddedIPythonCodeExecutor or LocalJupyterCodeExecutor not installed."
-)
 
+class DockerJupyterExecutor(JupyterCodeExecutor):
+    def __init__(self, **kwargs):
+        jupyter_server = DockerJupyterServer()
+        super().__init__(jupyter_server=jupyter_server, **kwargs)
 
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_is_code_executor(cls) -> None:
-    assert isinstance(cls, CodeExecutor)
 
+class LocalJupyterCodeExecutor(JupyterCodeExecutor):
+    def __init__(self, **kwargs):
+        jupyter_server = LocalJupyterServer()
+        super().__init__(jupyter_server=jupyter_server, **kwargs)
 
-@pytest.mark.skipif(skip, reason=skip_reason)
-def test_create_dict() -> None:
-    config: dict[str, Union[str, CodeExecutor]] = {"executor": "ipython-embedded"}
-    executor = CodeExecutorFactory.create(config)
-    assert isinstance(executor, EmbeddedIPythonCodeExecutor)
 
+# Skip on windows due to kernelgateway bug https://github.com/jupyter-server/kernel_gateway/issues/398
+if sys.platform == "win32":
+    classes_to_test = [EmbeddedIPythonCodeExecutor]
+else:
+    classes_to_test = [EmbeddedIPythonCodeExecutor, LocalJupyterCodeExecutor]
 
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_create(cls) -> None:
-    config = {"executor": cls()}
-    executor = CodeExecutorFactory.create(config)
-    assert executor is config["executor"]
-
+if not is_docker_running() or not decide_use_docker(use_docker=None):
+    skip_docker_test = True
+else:
+    skip_docker_test = False
+if not skip_docker_test:
+    classes_to_test.append(DockerJupyterExecutor)
+
+
+@skip_on_missing_imports(
+    [
+        "websocket",
+        "requests",
+        "jupyter_client",
+        "ipykernel",
+    ],
+    "jupyter_executor",
+)
+@skip_on_missing_jupyter_kernel_gateway()
+class TestCodeExecutor:
+    def test_import_utils(self) -> None:
+        pass
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_is_code_executor(self, cls) -> None:
+        assert isinstance(cls, CodeExecutor)
+
+    def test_create_dict(self) -> None:
+        config: dict[str, Union[str, CodeExecutor]] = {"executor": "ipython-embedded"}
+        executor = CodeExecutorFactory.create(config)
+        assert isinstance(executor, EmbeddedIPythonCodeExecutor)
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_create(self, cls) -> None:
+        config = {"executor": cls()}
+        executor = CodeExecutorFactory.create(config)
+        assert executor is config["executor"]
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_init(self, cls) -> None:
+        executor = cls(timeout=10, kernel_name="python3", output_dir=".")
+        assert executor._timeout == 10 and executor._kernel_name == "python3" and executor._output_dir == Path()
+
+        # Try invalid output directory.
+        with pytest.raises(ValueError, match="Output directory .* does not exist."):
+            executor = cls(timeout=111, kernel_name="python3", output_dir="/invalid/directory")
+
+        # Try invalid kernel name.
+        with pytest.raises(ValueError, match="Kernel .* is not installed."):
+            executor = cls(timeout=111, kernel_name="invalid_kernel_name", output_dir=".")
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_execute_code_single_code_block(self, cls) -> None:
+        executor = cls()
+        code_blocks = [CodeBlock(code="import sys\nprint('hello world!')", language="python")]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code == 0 and "hello world!" in code_result.output
 
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_init(cls) -> None:
-    executor = cls(timeout=10, kernel_name="python3", output_dir=".")
-    assert executor._timeout == 10 and executor._kernel_name == "python3" and executor._output_dir == Path()
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_execute_code_multiple_code_blocks(self, cls) -> None:
+        executor = cls()
+        code_blocks = [
+            CodeBlock(code="import sys\na = 123 + 123\n", language="python"),
+            CodeBlock(code="print(a)", language="python"),
+        ]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code == 0 and "246" in code_result.output
 
-    # Try invalid output directory.
-    with pytest.raises(ValueError, match="Output directory .* does not exist."):
-        executor = cls(timeout=111, kernel_name="python3", output_dir="/invalid/directory")
+        msg = """
+def test_function(a, b):
+    return a + b
+"""
+        code_blocks = [
+            CodeBlock(code=msg, language="python"),
+            CodeBlock(code="test_function(431, 423)", language="python"),
+        ]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code == 0 and "854" in code_result.output
 
-    # Try invalid kernel name.
-    with pytest.raises(ValueError, match="Kernel .* is not installed."):
-        executor = cls(timeout=111, kernel_name="invalid_kernel_name", output_dir=".")
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_execute_code_bash_script(self, cls) -> None:
+        executor = cls()
+        # Test bash script.
+        code_blocks = [CodeBlock(code='!echo "hello world!"', language="bash")]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code == 0 and "hello world!" in code_result.output
 
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_timeout(self, cls) -> None:
+        executor = cls(timeout=1)
+        code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code and "Timeout" in code_result.output
 
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_execute_code_single_code_block(cls) -> None:
-    executor = cls()
-    code_blocks = [CodeBlock(code="import sys\nprint('hello world!')", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and "hello world!" in code_result.output
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_silent_pip_install(self, cls) -> None:
+        executor = cls(timeout=600)
+        code_blocks = [CodeBlock(code="!pip install matplotlib numpy", language="python")]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code == 0 and code_result.output.strip() == ""
 
+        none_existing_package = uuid.uuid4().hex
+        code_blocks = [CodeBlock(code=f"!pip install matplotlib_{none_existing_package}", language="python")]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code == 0 and "ERROR: " in code_result.output
 
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_execute_code_multiple_code_blocks(cls) -> None:
-    executor = cls()
-    code_blocks = [
-        CodeBlock(code="import sys\na = 123 + 123\n", language="python"),
-        CodeBlock(code="print(a)", language="python"),
-    ]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and "246" in code_result.output
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_restart(self, cls) -> None:
+        executor = cls()
+        code_blocks = [CodeBlock(code="x = 123", language="python")]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code == 0 and code_result.output.strip() == ""
 
-    msg = """
-def test_function(a, b):
-    return a + b
-"""
-    code_blocks = [
-        CodeBlock(code=msg, language="python"),
-        CodeBlock(code="test_function(431, 423)", language="python"),
-    ]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and "854" in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_execute_code_bash_script(cls) -> None:
-    executor = cls()
-    # Test bash script.
-    code_blocks = [CodeBlock(code='!echo "hello world!"', language="bash")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and "hello world!" in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_timeout(cls) -> None:
-    executor = cls(timeout=1)
-    code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code and "Timeout" in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_silent_pip_install(cls) -> None:
-    executor = cls(timeout=600)
-    code_blocks = [CodeBlock(code="!pip install matplotlib numpy", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and code_result.output.strip() == ""
-
-    none_existing_package = uuid.uuid4().hex
-    code_blocks = [CodeBlock(code=f"!pip install matplotlib_{none_existing_package}", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and "ERROR: " in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_restart(cls) -> None:
-    executor = cls()
-    code_blocks = [CodeBlock(code="x = 123", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and code_result.output.strip() == ""
-
-    executor.restart()
-    code_blocks = [CodeBlock(code="print(x)", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code and "NameError" in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_save_image(cls) -> None:
-    with tempfile.TemporaryDirectory() as temp_dir:
-        executor = cls(output_dir=temp_dir)
-        # Install matplotlib.
-        code_blocks = [CodeBlock(code="!pip install matplotlib", language="python")]
+        executor.restart()
+        code_blocks = [CodeBlock(code="print(x)", language="python")]
+        code_result = executor.execute_code_blocks(code_blocks)
+        assert code_result.exit_code and "NameError" in code_result.output
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_save_image(self, cls) -> None:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            executor = cls(output_dir=temp_dir)
+            # Install matplotlib.
+            code_blocks = [CodeBlock(code="!pip install matplotlib", language="python")]
+            code_result = executor.execute_code_blocks(code_blocks)
+            assert code_result.exit_code == 0 and code_result.output.strip() == ""
+
+            # Test saving image.
+            code_blocks = [
+                CodeBlock(code="import matplotlib.pyplot as plt\nplt.plot([1, 2, 3, 4])\nplt.show()", language="python")
+            ]
+            code_result = executor.execute_code_blocks(code_blocks)
+            assert code_result.exit_code == 0
+            assert os.path.exists(code_result.output_files[0])
+            assert f"Image data saved to {code_result.output_files[0]}" in code_result.output
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_timeout_preserves_kernel_state(self, cls: type[CodeExecutor]) -> None:
+        executor = cls(timeout=1)
+        code_blocks = [CodeBlock(code="x = 123", language="python")]
         code_result = executor.execute_code_blocks(code_blocks)
         assert code_result.exit_code == 0 and code_result.output.strip() == ""
 
-        # Test saving image.
-        code_blocks = [
-            CodeBlock(code="import matplotlib.pyplot as plt\nplt.plot([1, 2, 3, 4])\nplt.show()", language="python")
-        ]
+        code_blocks = [CodeBlock(code="import time; time.sleep(10)", language="python")]
         code_result = executor.execute_code_blocks(code_blocks)
-        assert code_result.exit_code == 0
-        assert os.path.exists(code_result.output_files[0])
-        assert f"Image data saved to {code_result.output_files[0]}" in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_timeout_preserves_kernel_state(cls: type[CodeExecutor]) -> None:
-    executor = cls(timeout=1)
-    code_blocks = [CodeBlock(code="x = 123", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and code_result.output.strip() == ""
-
-    code_blocks = [CodeBlock(code="import time; time.sleep(10)", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code != 0 and "Timeout" in code_result.output
-
-    code_blocks = [CodeBlock(code="print(x)", language="python")]
-    code_result = executor.execute_code_blocks(code_blocks)
-    assert code_result.exit_code == 0 and "123" in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_save_html(cls) -> None:
-    with tempfile.TemporaryDirectory() as temp_dir:
-        executor = cls(output_dir=temp_dir)
-        # Test saving html.
-        code_blocks = [
-            CodeBlock(code="from IPython.display import HTML\nHTML('<h1>Hello, world!</h1>')", language="python")
-        ]
+        assert code_result.exit_code != 0 and "Timeout" in code_result.output
+
+        code_blocks = [CodeBlock(code="print(x)", language="python")]
         code_result = executor.execute_code_blocks(code_blocks)
-        assert code_result.exit_code == 0
-        assert os.path.exists(code_result.output_files[0])
-        assert f"HTML data saved to {code_result.output_files[0]}" in code_result.output
-
-
-@pytest.mark.skipif(skip, reason=skip_reason)
-@pytest.mark.parametrize("cls", classes_to_test)
-def test_conversable_agent_code_execution(cls) -> None:
-    agent = ConversableAgent(
-        "user_proxy",
-        llm_config=False,
-        code_execution_config={"executor": cls()},
-    )
-    msg = """
+        assert code_result.exit_code == 0 and "123" in code_result.output
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_save_html(self, cls) -> None:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            executor = cls(output_dir=temp_dir)
+            # Test saving html.
+            code_blocks = [
+                CodeBlock(code="from IPython.display import HTML\nHTML('<h1>Hello, world!</h1>')", language="python")
+            ]
+            code_result = executor.execute_code_blocks(code_blocks)
+            assert code_result.exit_code == 0
+            assert os.path.exists(code_result.output_files[0])
+            assert f"HTML data saved to {code_result.output_files[0]}" in code_result.output
+
+    @pytest.mark.parametrize("cls", classes_to_test)
+    def test_conversable_agent_code_execution(self, cls) -> None:
+        agent = ConversableAgent(
+            "user_proxy",
+            llm_config=False,
+            code_execution_config={"executor": cls()},
+        )
+        msg = """
 Run this code:
 ```python
 def test_function(a, b):
@@ -252,10 +238,10 @@ def test_function(a, b):
 print(test_function(123, 4))
 ```
 """
-    with pytest.MonkeyPatch.context() as mp:
-        mp.setenv("OPENAI_API_KEY", MOCK_OPEN_AI_API_KEY)
-        reply = agent.generate_reply(
-            [{"role": "user", "content": msg}],
-            sender=ConversableAgent("user", llm_config=False, code_execution_config=False),
-        )
-        assert "492" in reply  # type: ignore[operator]
+        with pytest.MonkeyPatch.context() as mp:
+            mp.setenv("OPENAI_API_KEY", MOCK_OPEN_AI_API_KEY)
+            reply = agent.generate_reply(
+                [{"role": "user", "content": msg}],
+                sender=ConversableAgent("user", llm_config=False, code_execution_config=False),
+            )
+            assert "492" in reply  # type: ignore[operator]
diff --git a/test/oai/test_anthropic.py b/test/oai/test_anthropic.py
index 4b15cd32d..fafe908fb 100644
--- a/test/oai/test_anthropic.py
+++ b/test/oai/test_anthropic.py
@@ -11,11 +11,11 @@
 import pytest
 
 from autogen.import_utils import optional_import_block
+from autogen.oai.anthropic import AnthropicClient, _calculate_cost
 
 with optional_import_block() as result:
     from anthropic.types import Message, TextBlock
 
-    from autogen.oai.anthropic import AnthropicClient, _calculate_cost
 
 skip = not result.is_successful
 
diff --git a/test/oai/test_cohere.py b/test/oai/test_cohere.py
index bdb8b2c79..e4db3692e 100644
--- a/test/oai/test_cohere.py
+++ b/test/oai/test_cohere.py
@@ -11,9 +11,10 @@
 import pytest
 
 from autogen.import_utils import optional_import_block
+from autogen.oai.cohere import CohereClient, calculate_cohere_cost
 
 with optional_import_block() as result:
-    from autogen.oai.cohere import CohereClient, calculate_cohere_cost
+    from cohere import Client as Cohere  # noqa: F401
 
 skip = not result.is_successful
 
diff --git a/test/oai/test_gemini.py b/test/oai/test_gemini.py
index 27d6f3292..71a7cca34 100644
--- a/test/oai/test_gemini.py
+++ b/test/oai/test_gemini.py
@@ -13,9 +13,13 @@
 from pydantic import BaseModel
 
 from autogen.import_utils import optional_import_block
+from autogen.oai.gemini import GeminiClient
 
 with optional_import_block() as result:
+    import google.ai  # noqa: F401
     import google.auth  # noqa: F401
+    import vertexai  # noqa: F401
+    from PIL import Image  # noqa: F401
     from google.api_core.exceptions import InternalServerError
     from google.auth.credentials import Credentials
     from google.cloud.aiplatform.initializer import global_config as vertexai_global_config
@@ -25,8 +29,6 @@
     from vertexai.generative_models import HarmCategory as VertexAIHarmCategory
     from vertexai.generative_models import SafetySetting as VertexAISafetySetting
 
-    from autogen.oai.gemini import GeminiClient
-
 skip = not result.is_successful
 
 
diff --git a/test/oai/test_groq.py b/test/oai/test_groq.py
index 396040816..2519cee0e 100644
--- a/test/oai/test_groq.py
+++ b/test/oai/test_groq.py
@@ -9,9 +9,10 @@
 import pytest
 
 from autogen.import_utils import optional_import_block
+from autogen.oai.groq import GroqClient, calculate_groq_cost
 
 with optional_import_block() as result:
-    from autogen.oai.groq import GroqClient, calculate_groq_cost
+    from groq import Groq  # noqa: F401
 
 skip = not result.is_successful
 
diff --git a/test/oai/test_mistral.py b/test/oai/test_mistral.py
index 2aeac48c9..c8428d42c 100644
--- a/test/oai/test_mistral.py
+++ b/test/oai/test_mistral.py
@@ -9,6 +9,7 @@
 import pytest
 
 from autogen.import_utils import optional_import_block
+from autogen.oai.mistral import MistralAIClient, calculate_mistral_cost
 
 with optional_import_block() as result:
     from mistralai import (
@@ -22,7 +23,6 @@
         UserMessage,  # noqa: F401
     )
 
-    from autogen.oai.mistral import MistralAIClient, calculate_mistral_cost
 
 skip = not result.is_successful
 
diff --git a/test/test_import.py b/test/test_import.py
new file mode 100644
index 000000000..e77a0885e
--- /dev/null
+++ b/test/test_import.py
@@ -0,0 +1,84 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+import importlib
+import pkgutil
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Optional
+
+import pytest
+
+
+@contextmanager
+def add_to_sys_path(path: Optional[Path]) -> Iterator[None]:
+    if path is None:
+        yield
+        return
+
+    if not path.exists():
+        raise ValueError(f"Path {path} does not exist")
+
+    sys.path.append(str(path))
+    try:
+        yield
+    finally:
+        sys.path.remove(str(path))
+
+
+def list_submodules(module_name: str, *, include_path: Optional[Path] = None, include_root: bool = True) -> list[str]:
+    """List all submodules of a given module.
+
+    Args:
+        module_name (str): The name of the module to list submodules for.
+        include_path (Optional[Path], optional): The path to the module. Defaults to None.
+        include_root (bool, optional): Whether to include the root module in the list. Defaults to True.
+
+    Returns:
+        list: A list of submodule names.
+    """
+    with add_to_sys_path(include_path):
+        try:
+            module = importlib.import_module(module_name)  # nosemgrep
+        except Exception:
+            return []
+
+        # Get the path of the module. This is necessary to find its submodules.
+        module_path = module.__path__
+
+        # Initialize an empty list to store the names of submodules
+        submodules = [module_name] if include_root else []
+
+        # Iterate over the submodules in the module's path
+        for _, name, ispkg in pkgutil.iter_modules(module_path, prefix=f"{module_name}."):
+            # Add the name of each submodule to the list
+            submodules.append(name)
+
+            if ispkg:
+                submodules.extend(list_submodules(name, include_root=False))
+
+        # Return the list of submodule names
+        return submodules
+
+
+def test_list_submodules() -> None:
+    # Specify the name of the module you want to inspect
+    module_name = "autogen"
+
+    # Get the list of submodules for the specified module
+    submodules = list_submodules(module_name)
+
+    assert len(submodules) > 0
+    assert "autogen" in submodules
+    assert "autogen.io" in submodules
+    assert "autogen.coding.jupyter" in submodules
+
+
+# todo: we should always run this
+@pytest.mark.parametrize("module", list_submodules("autogen"))
+def test_submodules(module: str) -> None:
+    importlib.import_module(module)  # nosemgrep
diff --git a/test/test_import_utils.py b/test/test_import_utils.py
index fa08f1318..89d4c787e 100644
--- a/test/test_import_utils.py
+++ b/test/test_import_utils.py
@@ -10,7 +10,7 @@
 
 
 class TestOptionalImportBlock:
-    def test_optional_import_block(self):
+    def test_optional_import_block(self) -> None:
         with optional_import_block():
             import ast
 
@@ -32,8 +32,8 @@ def test_optional_import_block(self):
 
 
 class TestRequiresOptionalImportCallables:
-    def test_function_attributes(self):
-        def dummy_function():
+    def test_function_attributes(self) -> None:
+        def dummy_function() -> None:
             """Dummy function to test requires_optional_import"""
             pass
 
@@ -52,9 +52,9 @@ def dummy_function():
         ):
             actual()
 
-    def test_function_call(self):
+    def test_function_call(self) -> None:
         @require_optional_import("some_optional_module", "optional_dep")
-        def dummy_function():
+        def dummy_function() -> None:
             """Dummy function to test requires_optional_import"""
             pass
 
@@ -64,9 +64,9 @@ def dummy_function():
         ):
             dummy_function()
 
-    def test_method_attributes(self):
+    def test_method_attributes(self) -> None:
         class DummyClass:
-            def dummy_method(self):
+            def dummy_method(self) -> None:
                 """Dummy method to test requires_optional_import"""
                 pass
 
@@ -76,7 +76,7 @@ def dummy_method(self):
         DummyClass.__module__ = "some_random_module.dummy_stuff"
         DummyClass.dummy_method.__module__ = "some_random_module.dummy_stuff"
 
-        DummyClass.dummy_method = require_optional_import("some_optional_module", "optional_dep")(
+        DummyClass.dummy_method = require_optional_import("some_optional_module", "optional_dep")(  # type: ignore[method-assign]
             DummyClass.dummy_method
         )
 
@@ -92,10 +92,10 @@ def dummy_method(self):
         ):
             dummy.dummy_method()
 
-    def test_method_call(self):
+    def test_method_call(self) -> None:
         class DummyClass:
             @require_optional_import("some_optional_module", "optional_dep")
-            def dummy_method(self):
+            def dummy_method(self) -> None:
                 """Dummy method to test requires_optional_import"""
                 pass
 
@@ -106,11 +106,11 @@ def dummy_method(self):
         ):
             dummy.dummy_method()
 
-    def test_static_call(self):
+    def test_static_call(self) -> None:
         class DummyClass:
             @require_optional_import("some_optional_module", "optional_dep")
             @staticmethod
-            def dummy_static_function():
+            def dummy_static_function() -> None:
                 """Dummy static function to test requires_optional_import"""
                 pass
 
@@ -121,10 +121,10 @@ def dummy_static_function():
         ):
             dummy.dummy_static_function()
 
-    def test_property_call(self):
+    def test_property_call(self) -> None:
         class DummyClass:
-            @require_optional_import("some_optional_module", "optional_dep")
             @property
+            @require_optional_import("some_optional_module", "optional_dep")
             def dummy_property(self) -> int:
                 """Dummy property to test requires_optional_import"""
                 return 4
@@ -142,17 +142,17 @@ class TestRequiresOptionalImportClasses:
     def dummy_cls(self) -> Type[Any]:
         @require_optional_import("some_optional_module", "optional_dep")
         class DummyClass:
-            def dummy_method(self):
+            def dummy_method(self) -> None:
                 """Dummy method to test requires_optional_import"""
                 pass
 
             @staticmethod
-            def dummy_static_method():
+            def dummy_static_method() -> None:
                 """Dummy static method to test requires_optional_import"""
                 pass
 
             @classmethod
-            def dummy_class_method(cls):
+            def dummy_class_method(cls) -> None:
                 """Dummy class method to test requires_optional_import"""
                 pass
 
@@ -163,7 +163,7 @@ def dummy_property(self) -> int:
 
         return DummyClass
 
-    def test_class_init_call(self, dummy_cls: Type[Any]):
+    def test_class_init_call(self, dummy_cls: Type[Any]) -> None:
         with pytest.raises(
             ImportError,
             match=r"Module 'some_optional_module' needed for __init__ is missing, please install it using 'pip install ag2\[optional_dep\]'",