Use skip_on_missing_imports in test files to mark tests (#633)

* Use skip_on_missing_imports in test files to mark tests

* Use skip_on_missing_imports to mark tests in test/oai files

* Use skip_on_missing_imports to mark tests in test/io files

* Use skip_on_missing_imports to mark tests in test/interop files

* Use skip_on_missing_imports to mark tests in test/coding files

* Use skip_on_missing_imports to mark tests in test/cache files

* Use skip_on_missing_imports to mark tests in test/agentchat files

* Use skip_on_missing_imports to mark tests in test/agentchat/contrib files

* Use skip_on_missing_imports to mark tests in test/agentchat/contrib/vectordb files

* Use skip_on_missing_imports to mark tests in test/agentchat/contrib/retrievechat files

* Fix bug in test

* Use skip_on_missing_imports to mark tests in test/agentchat/contrib/graph_rag files

* Use skip_on_missing_imports to mark tests in test/agentchat/contrib/capabilities files

autogen/import_utils.py

@@ -268,7 +268,7 @@ def decorator(o: T) -> T:
     return decorator
 
 
-def skip_on_missing_imports(modules: Union[str, Iterable[str]], dep_target: str) -> Callable[[T], T]:
+def skip_on_missing_imports(modules: Union[str, Iterable[str]], dep_target: Optional[str] = None) -> Callable[[T], T]:
     """Decorator to skip a test if an optional module is missing
 
     Args:
@@ -287,8 +287,9 @@ def decorator(o: T) -> T:
         def decorator(o: T) -> T:
             import pytest
 
+            install_target = "" if dep_target is None else f"[{dep_target}]"
             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}]'"
+                f"Missing module{'s' if len(missing_modules) > 1 else ''}: {', '.join(missing_modules)}. Install using 'pip install ag2{install_target}'"
             )(o)
 
     return decorator

test/agentchat/contrib/capabilities/test_image_generation_capability.py

@@ -25,7 +25,6 @@
 with optional_import_block() as result:
     from PIL import Image
 
-skip_requirement = not result.is_successful
 
 filter_dict = {"model": ["gpt-4o-mini"]}
 
@@ -102,7 +101,7 @@ def test_dalle_image_generator(dalle_config: dict[str, Any]):
 # Using cartesian product to generate all possible combinations of resolution, quality, and prompt
 @pytest.mark.parametrize("gen_config_1", itertools.product(RESOLUTIONS, QUALITIES, PROMPTS))
 @pytest.mark.parametrize("gen_config_2", itertools.product(RESOLUTIONS, QUALITIES, PROMPTS))
-@pytest.mark.skipif(skip_requirement, reason="Dependencies are not installed.")
+@skip_on_missing_imports(["PIL"], "unknown")
 def test_dalle_image_generator_cache_key(
     dalle_config: dict[str, Any], gen_config_1: tuple[str, str, str], gen_config_2: tuple[str, str, str]
 ):
@@ -125,7 +124,7 @@ def test_dalle_image_generator_cache_key(
         assert cache_key_1 != cache_key_2
 
 
-@pytest.mark.skipif(skip_requirement, reason="Dependencies are not installed.")
+@skip_on_missing_imports(["PIL"], "unknown")
 def test_image_generation_capability_positive(monkeypatch, image_gen_capability):
     """Tests ImageGeneration capability to generate images by calling the ImageGenerator.
 
@@ -153,7 +152,7 @@ def test_image_generation_capability_positive(monkeypatch, image_gen_capability)
     assert auto_reply not in processed_message
 
 
-@pytest.mark.skipif(skip_requirement, reason="Dependencies are not installed.")
+@skip_on_missing_imports(["PIL"], "unknown")
 def test_image_generation_capability_negative(monkeypatch, image_gen_capability):
     """Tests ImageGeneration capability to generate images by calling the ImageGenerator.
 
@@ -181,7 +180,7 @@ def test_image_generation_capability_negative(monkeypatch, image_gen_capability)
     assert auto_reply == processed_message
 
 
-@pytest.mark.skipif(skip_requirement, reason="Dependencies are not installed.")
+@skip_on_missing_imports(["PIL"], "unknown")
 def test_image_generation_capability_cache(monkeypatch):
     """Tests ImageGeneration capability to cache the generated images."""
     test_image_size = (256, 256)

test/agentchat/contrib/capabilities/test_teachable_agent.py

@@ -11,17 +11,10 @@
 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 autogen.import_utils import skip_on_missing_imports
 
 from ....conftest import Credentials
 
-with optional_import_block() as result:
-    import chromadb  # noqa: F401
-
-
-skip = not result.is_successful
-
-
 # Specify the model to use by uncommenting one of the following lines.
 # filter_dict={"model": ["gpt-4-1106-preview"]}
 # filter_dict={"model": ["gpt-4-0613"]}
@@ -130,10 +123,7 @@ def use_task_advice_pair_phrasing(credentials: Credentials):
 
 
 @pytest.mark.openai
-@pytest.mark.skipif(
-    skip,
-    reason="do not run if dependency is not installed or requested to skip",
-)
+@skip_on_missing_imports(["chromadb"], "teachable")
 def test_teachability_code_paths(credentials_gpt_4o_mini: Credentials):
     """Runs this file's unit tests."""
     total_num_errors, total_num_tests = 0, 0
@@ -162,10 +152,7 @@ def test_teachability_code_paths(credentials_gpt_4o_mini: Credentials):
 
 
 @pytest.mark.openai
-@pytest.mark.skipif(
-    skip,
-    reason="do not run if dependency is not installed or requested to skip",
-)
+@skip_on_missing_imports(["chromadb"], "teachable")
 def test_teachability_accuracy(credentials_gpt_4o_mini: Credentials):
     """A very cheap and fast test of teachability accuracy."""
     print(colored("\nTEST TEACHABILITY ACCURACY", "light_cyan"))

test/agentchat/contrib/capabilities/test_vision_capability.py

@@ -11,15 +11,12 @@
 
 from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
 from autogen.agentchat.conversable_agent import ConversableAgent
-from autogen.import_utils import optional_import_block
+from autogen.import_utils import optional_import_block, skip_on_missing_imports
 
 with optional_import_block() as result:
     from PIL import Image  # noqa: F401
 
 
-skip_test = not result.is_successful
-
-
 @pytest.fixture
 def lmm_config():
     return {
@@ -42,24 +39,14 @@ def conversable_agent():
     return ConversableAgent(name="conversable_agent", llm_config=False)
 
 
-@pytest.mark.skipif(
-    skip_test,
-    reason="do not run if dependency is not installed",
-)
+@skip_on_missing_imports(["PIL"], "unknown")
 def test_add_to_conversable_agent(vision_capability, conversable_agent):
     vision_capability.add_to_agent(conversable_agent)
     assert hasattr(conversable_agent, "process_last_received_message")
 
 
-@pytest.mark.skipif(
-    skip_test,
-    reason="do not run if dependency is not installed",
-)
+@skip_on_missing_imports(["PIL"], "unknown")
 @patch("autogen.oai.client.OpenAIWrapper")
-@pytest.mark.skipif(
-    skip_test,
-    reason="do not run if dependency is not installed",
-)
 def test_process_last_received_message_text(mock_lmm_client, vision_capability):
     mock_lmm_client.create.return_value = MagicMock(choices=[MagicMock(message=MagicMock(content="A description"))])
     content = "Test message without image"
@@ -76,10 +63,7 @@ def test_process_last_received_message_text(mock_lmm_client, vision_capability):
     "autogen.agentchat.contrib.capabilities.vision_capability.VisionCapability._get_image_caption",
     return_value="A sample image caption.",
 )
-@pytest.mark.skipif(
-    skip_test,
-    reason="do not run if dependency is not installed",
-)
+@skip_on_missing_imports(["PIL"], "unknown")
 def test_process_last_received_message_with_image(
     mock_get_caption, mock_convert_base64, mock_get_image_data, vision_capability
 ):
@@ -105,10 +89,7 @@ def caption_func(image_url: str, image_data=None, lmm_client=None) -> str:
     return caption_func
 
 
-@pytest.mark.skipif(
-    skip_test,
-    reason="do not run if dependency is not installed",
-)
+@skip_on_missing_imports(["PIL"], "unknown")
 class TestCustomCaptionFunc:
     def test_custom_caption_func_with_valid_url(self, custom_caption_func):
         """Test custom caption function with a valid image URL."""

test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py

@@ -13,22 +13,21 @@
     FalkorGraphQueryEngine,
     GraphStoreQueryResult,
 )
-from autogen.import_utils import optional_import_block
+from autogen.import_utils import optional_import_block, skip_on_missing_imports
 
 with optional_import_block() as result:
-    import falkordb  # noqa: F401
     from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation
 
-skip = not result.is_successful
 
 reason = "do not run on MacOS or windows OR dependency is not installed"
 
 
 @pytest.mark.openai
 @pytest.mark.skipif(
-    sys.platform in ["darwin", "win32"] or skip,
+    sys.platform in ["darwin", "win32"],
     reason=reason,
 )
+@skip_on_missing_imports(["falkordb", "graphrag_sdk"], "neo4j")
 def test_falkor_db_query_engine():
     """Test FalkorDB Query Engine.
     1. create a test FalkorDB Query Engine with a schema.

test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py

@@ -12,18 +12,10 @@
     GraphStoreQueryResult,
     Neo4jNativeGraphQueryEngine,
 )
-from autogen.import_utils import optional_import_block
+from autogen.import_utils import skip_on_missing_imports
 
 from ....conftest import reason
 
-with optional_import_block() as result:
-    from neo4j import GraphDatabase  # noqa: F401
-    from neo4j_graphrag.embeddings import Embedder  # noqa: F401
-
-
-skip = not result.is_successful
-
-
 # Configure the logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -115,9 +107,10 @@ def neo4j_native_query_engine_auto():
 
 @pytest.mark.openai
 @pytest.mark.skipif(
-    sys.platform in ["darwin", "win32"] or skip,
+    sys.platform in ["darwin", "win32"],
     reason=reason,
 )
+@skip_on_missing_imports(["neo4j", "neo4j_graphrag"], "neo4j")
 def test_neo4j_native_query_engine(neo4j_native_query_engine):
     """Test querying with initialized knowledge graph"""
     question = "Which company is the employer?"
@@ -129,9 +122,10 @@ def test_neo4j_native_query_engine(neo4j_native_query_engine):
 
 @pytest.mark.openai
 @pytest.mark.skipif(
-    sys.platform in ["darwin", "win32"] or skip,
+    sys.platform in ["darwin", "win32"],
     reason=reason,
 )
+@skip_on_missing_imports(["neo4j", "neo4j_graphrag"], "neo4j")
 def test_neo4j_native_query_auto(neo4j_native_query_engine_auto):
     """Test querying with auto-generated property graph"""
     question = "Which company is the employer?"

test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py

@@ -15,17 +15,10 @@
     GraphStoreQueryResult,
     Neo4jGraphQueryEngine,
 )
-from autogen.import_utils import optional_import_block
+from autogen.import_utils import skip_on_missing_imports
 
 from ....conftest import reason
 
-with optional_import_block() as result:
-    from llama_index.core import PropertyGraphIndex  # noqa: F401
-
-
-skip = not result.is_successful
-
-
 # Configure the logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -35,6 +28,7 @@
 
 # Test fixture for creating and initializing a query engine with a JSON input file
 @pytest.fixture(scope="module")
+@skip_on_missing_imports(["llama_index"], "neo4j")
 def neo4j_query_engine_with_json():
     input_path = "./test/agentchat/contrib/graph_rag/layout_parser_paper_parsed_elements.json"
     input_documents = [Document(doctype=DocumentType.JSON, path_or_url=input_path)]
@@ -122,9 +116,10 @@ def neo4j_query_engine_auto():
 
 @pytest.mark.openai
 @pytest.mark.skipif(
-    sys.platform in ["darwin", "win32"] or skip,
+    sys.platform in ["darwin", "win32"],
     reason=reason,
 )
+@skip_on_missing_imports(["llama_index"], "neo4j")
 def test_neo4j_query_engine(neo4j_query_engine):
     """Test querying functionality of the Neo4j Query Engine."""
     question = "Which company is the employer?"
@@ -139,9 +134,10 @@ def test_neo4j_query_engine(neo4j_query_engine):
 
 @pytest.mark.openai
 @pytest.mark.skipif(
-    sys.platform in ["darwin", "win32"] or skip,
+    sys.platform in ["darwin", "win32"],
     reason=reason,
 )
+@skip_on_missing_imports(["llama_index"], "neo4j")
 def test_neo4j_add_records(neo4j_query_engine):
     """Test the add_records functionality of the Neo4j Query Engine."""
     input_path = "./test/agentchat/contrib/graph_rag/the_matrix.txt"
@@ -161,9 +157,10 @@ def test_neo4j_add_records(neo4j_query_engine):
 
 @pytest.mark.openai
 @pytest.mark.skipif(
-    sys.platform in ["darwin", "win32"] or skip,
+    sys.platform in ["darwin", "win32"],
     reason=reason,
 )
+@skip_on_missing_imports(["llama_index"], "neo4j")
 def test_neo4j_auto(neo4j_query_engine_auto):
     """Test querying with auto-generated property graph"""
     question = "Which company is the employer?"
@@ -175,9 +172,10 @@ def test_neo4j_auto(neo4j_query_engine_auto):
 
 @pytest.mark.openai
 @pytest.mark.skipif(
-    sys.platform in ["darwin", "win32"] or skip,
+    sys.platform in ["darwin", "win32"],
     reason=reason,
 )
+@skip_on_missing_imports(["llama_index"], "neo4j")
 def test_neo4j_json_auto(neo4j_query_engine_with_json):
     """Test querying with auto-generated property graph from a JSON file."""
     question = "What are current layout detection models in the LayoutParser model zoo?"

test/agentchat/contrib/retrievechat/test_pgvector_retrievechat.py

@@ -9,32 +9,25 @@
 import os
 
 import pytest
-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 autogen.import_utils import optional_import_block, skip_on_missing_imports
 
 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 sentence_transformers import SentenceTransformer
 
 
-skip = not result.is_successful
-
 test_dir = os.path.join(os.path.dirname(__file__), "../../..", "test_files")
 
 
 @pytest.mark.openai
-@pytest.mark.skipif(
-    skip,
-    reason="dependency is not installed OR requested to skip",
-)
+@skip_on_missing_imports(["chromadb", "pgvector", "IPython", "sentence_transformers"], "retrievechat-pgvector")
 def test_retrievechat(credentials_gpt_4o_mini: Credentials):
     conversations = {}