all multiple models to pytest --vlm-model-id

langchain-ai · Aug 22, 2024 · bb2fa1d · bb2fa1d
1 parent 63b9f0e
commit bb2fa1d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/libs/ai-endpoints/tests/integration_tests/conftest.py b/libs/ai-endpoints/tests/integration_tests/conftest.py
@@ -54,7 +54,8 @@ def pytest_addoption(parser: pytest.Parser) -> None:
     parser.addoption(
         "--vlm-model-id",
         action="store",
-        help="Run tests for a specific vlm model",
+        nargs="+",
+        help="Run tests for a specific vlm model or list of models",
     )
     parser.addoption(
         "--all-models",
@@ -120,8 +121,8 @@ def get_all_known_models() -> List[Model]:
 
     if "vlm_model" in metafunc.fixturenames:
         models = ["nvidia/neva-22b"]
-        if model := metafunc.config.getoption("vlm_model_id"):
-            models = [model]
+        if model_list := metafunc.config.getoption("vlm_model_id"):
+            models = model_list
         if metafunc.config.getoption("all_models"):
             models = [
                 model.id