Use optional_import_block in agentchat/contrib tests

test/agentchat/contrib/capabilities/test_image_generation_capability.py

@@ -15,19 +15,18 @@
 from autogen.agentchat.conversable_agent import ConversableAgent
 from autogen.agentchat.user_proxy_agent import UserProxyAgent
 from autogen.cache.cache import Cache
+from autogen.import_utils import optional_import_block
 from autogen.oai import openai_utils
 
-try:
+from ....conftest import MOCK_OPEN_AI_API_KEY
+
+with optional_import_block() as result:
     from PIL import Image
 
     from autogen.agentchat.contrib.capabilities import generate_images
     from autogen.agentchat.contrib.img_utils import get_pil_image
-except ImportError:
-    skip_requirement = True
-else:
-    skip_requirement = False
 
-from ....conftest import MOCK_OPEN_AI_API_KEY
+skip_requirement = not result.is_successful
 
 filter_dict = {"model": ["gpt-4o-mini"]}
 

test/agentchat/contrib/capabilities/test_teachable_agent.py

@@ -10,15 +10,14 @@
 
 from autogen import ConversableAgent
 from autogen.formatting_utils import colored
+from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
 
-try:
+with optional_import_block() as result:
     from autogen.agentchat.contrib.capabilities.teachability import Teachability
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
 
 
 # Specify the model to use by uncommenting one of the following lines.

test/agentchat/contrib/capabilities/test_transforms.py

@@ -17,6 +17,7 @@
     TextMessageContentName,
 )
 from autogen.agentchat.contrib.capabilities.transforms_util import count_text_tokens
+from autogen.import_utils import optional_import_block
 
 
 class _MockTextCompressor:
@@ -104,12 +105,11 @@ def get_messages_with_names_post_filtered() -> list[dict]:
 
 def get_text_compressors() -> list[TextCompressor]:
     compressors: list[TextCompressor] = [_MockTextCompressor()]
-    try:
+    with optional_import_block() as result:
         from autogen.agentchat.contrib.capabilities.text_compressors import LLMLingua
 
+    if result.is_successful:
         compressors.append(LLMLingua())
-    except ImportError:
-        pass
 
     return compressors
 

test/agentchat/contrib/capabilities/test_vision_capability.py

@@ -10,15 +10,14 @@
 import pytest
 
 from autogen.agentchat.conversable_agent import ConversableAgent
+from autogen.import_utils import optional_import_block
 
-try:
+with optional_import_block() as result:
     from PIL import Image  # noqa: F401
 
     from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
-except ImportError:
-    skip_test = True
-else:
-    skip_test = False
+
+skip_test = not result.is_successful
 
 
 @pytest.fixture

test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py

@@ -9,16 +9,16 @@
 import pytest
 from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation
 
-try:
+from autogen.import_utils import optional_import_block
+
+with optional_import_block() as result:
     from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
     from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import (
         FalkorGraphQueryEngine,
         GraphStoreQueryResult,
     )
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
 
 reason = "do not run on MacOS or windows OR dependency is not installed"
 

test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py

@@ -7,19 +7,20 @@
 
 import pytest
 
+from autogen.import_utils import optional_import_block
+
 from ....conftest import reason
 
-try:
+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,
     )
 
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
+
 
 # Configure the logging
 logging.basicConfig(level=logging.INFO)

test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py

@@ -10,19 +10,20 @@
 
 import pytest
 
+from autogen.import_utils import optional_import_block
+
 from ....conftest import reason
 
-try:
+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,
     )
 
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
+
 
 # Configure the logging
 logging.basicConfig(level=logging.INFO)

test/agentchat/contrib/retrievechat/test_pgvector_retrievechat.py

@@ -12,20 +12,18 @@
 from sentence_transformers import SentenceTransformer
 
 from autogen import AssistantAgent
+from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
 
-try:
+with optional_import_block() as result:
     import pgvector  # noqa: F401
 
     from autogen.agentchat.contrib.retrieve_user_proxy_agent import (
         RetrieveUserProxyAgent,
     )
-except ImportError:
-    skip = True
-else:
-    skip = False
 
+skip = not result.is_successful
 
 test_dir = os.path.join(os.path.dirname(__file__), "../../..", "test_files")
 

test/agentchat/contrib/retrievechat/test_qdrant_retrievechat.py

@@ -12,10 +12,11 @@
 import pytest
 
 from autogen import AssistantAgent
+from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
 
-try:
+with optional_import_block() as result:
     import fastembed  # noqa: F401
     from qdrant_client import QdrantClient
 
@@ -25,16 +26,12 @@
         query_qdrant,
     )
 
-    QDRANT_INSTALLED = True
-except ImportError:
-    QDRANT_INSTALLED = False
+QDRANT_INSTALLED = result.is_successful
 
-try:
+with optional_import_block() as result:
     import openai  # noqa: F401
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
 
 
 @pytest.mark.openai

test/agentchat/contrib/retrievechat/test_retrievechat.py

@@ -10,9 +10,11 @@
 
 import pytest
 
+from autogen.import_utils import optional_import_block
+
 from ....conftest import Credentials, reason
 
-try:
+with optional_import_block() as result:
     import chromadb
     import openai  # noqa: F401
     from chromadb.utils import embedding_functions as ef
@@ -21,10 +23,8 @@
     from autogen.agentchat.contrib.retrieve_user_proxy_agent import (
         RetrieveUserProxyAgent,
     )
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
 
 reason = "do not run on MacOS or windows OR dependency is not installed OR " + reason
 

test/agentchat/contrib/vectordb/test_chromadb.py

@@ -9,18 +9,18 @@
 
 import pytest
 
+from autogen.import_utils import optional_import_block
+
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 
-try:
+with optional_import_block() as result:
     import chromadb
     import chromadb.errors
     import sentence_transformers  # noqa: F401
 
     from autogen.agentchat.contrib.vectordb.chromadb import ChromaVectorDB
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
 
 
 @pytest.mark.skipif(skip, reason="dependency is not installed")

test/agentchat/contrib/vectordb/test_mongodb.py

@@ -12,13 +12,15 @@
 import pytest
 
 from autogen.agentchat.contrib.vectordb.base import Document
+from autogen.import_utils import optional_import_block
 
-try:
+with optional_import_block() as result:
     import pymongo  # noqa: F401
     import sentence_transformers  # noqa: F401
 
     from autogen.agentchat.contrib.vectordb.mongodb import MongoDBAtlasVectorDB
-except ImportError:
+
+if not result.is_successful:
     # To display warning in pyproject.toml [tool.pytest.ini_options] set log_cli = true
     logger = logging.getLogger(__name__)
     logger.warning(f"skipping {__name__}. It requires one to pip install pymongo or the extra [retrievechat-mongodb]")

test/agentchat/contrib/vectordb/test_pgvectordb.py

@@ -10,18 +10,18 @@
 
 import pytest
 
+from autogen.import_utils import optional_import_block
+
 from ....conftest import reason
 
-try:
+with optional_import_block() as result:
     import pgvector  # noqa: F401
     import psycopg
     import sentence_transformers  # noqa: F401
 
     from autogen.agentchat.contrib.vectordb.pgvectordb import PGVectorDB
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
 
 reason = "do not run on MacOS or windows OR dependency is not installed OR " + reason
 

test/agentchat/contrib/vectordb/test_qdrant.py

@@ -9,18 +9,18 @@
 
 import pytest
 
+from autogen.import_utils import optional_import_block
+
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 
-try:
+with optional_import_block() as result:
     import uuid
 
     from qdrant_client import QdrantClient
 
     from autogen.agentchat.contrib.vectordb.qdrant import QdrantVectorDB
-except ImportError:
-    skip = True
-else:
-    skip = False
+
+skip = not result.is_successful
 
 
 @pytest.mark.skipif(skip, reason="dependency is not installed")