speed up test execution (#619)

langchain-ai · Nov 27, 2024 · 7df7cad · 7df7cad
1 parent ef44612
commit 7df7cad
Show file tree

Hide file tree

Showing 8 changed files with 612 additions and 635 deletions.
diff --git a/libs/vertexai/langchain_google_vertexai/_anthropic_utils.py b/libs/vertexai/langchain_google_vertexai/_anthropic_utils.py
@@ -31,7 +31,7 @@
 from pydantic import BaseModel
 
 if TYPE_CHECKING:
-    from anthropic.types import RawMessageStreamEvent
+    from anthropic.types import RawMessageStreamEvent  # type: ignore[import-not-found]
 
 _message_type_lookups = {
     "human": "user",

diff --git a/libs/vertexai/langchain_google_vertexai/model_garden.py b/libs/vertexai/langchain_google_vertexai/model_garden.py
@@ -148,7 +148,7 @@ def __init__(self, **kwargs: Any) -> None:
 
     @model_validator(mode="after")
     def validate_environment(self) -> Self:
-        from anthropic import (
+        from anthropic import (  # type: ignore[import-not-found]
             AnthropicVertex,
             AsyncAnthropicVertex,
         )

diff --git a/libs/vertexai/poetry.lock b/libs/vertexai/poetry.lock
diff --git a/libs/vertexai/pyproject.toml b/libs/vertexai/pyproject.toml
@@ -20,13 +20,15 @@ anthropic = { extras = ["vertexai"], version = ">=0.35.0,<1", optional = true }
 langchain-mistralai = { version = ">=0.2.0,<1", optional = true }
 httpx = "^0.27.0"
 httpx-sse = "^0.4.0"
-pydantic = ">=2,<3"
+pydantic = ">=2.9,<2.10"
 
 [tool.poetry.group.test]
 optional = true
 
 [tool.poetry.group.test.dependencies]
 pytest = "^7.3.0"
+pytest-xdist = "^3.0"
+pytest-order =  "^1.3"
 freezegun = "^1.2.2"
 pytest-mock = "^3.10.0"
 syrupy = "^4.0.2"
@@ -81,6 +83,7 @@ mypy = "^1"
 types-google-cloud-ndb = "^2.2.0.20240106"
 types-protobuf = "^4.24.0.4"
 types-requests = "^2.31.0"
+anthropic = { extras = ["vertexai"], version = ">=0.35.0,<1", optional = true }
 
 
 [tool.poetry.group.dev.dependencies]

diff --git a/libs/vertexai/tests/conftest.py b/libs/vertexai/tests/conftest.py
@@ -11,7 +11,7 @@
 _LONG_FLAG = "long"
 _EXTENDED_FLAG = "extended"
 
-_PYTEST_FLAGS = [_RELEASE_FLAG, _GPU_FLAG, _LONG_FLAG, _EXTENDED_FLAG]
+_PYTEST_FLAGS = [_RELEASE_FLAG, _GPU_FLAG, _LONG_FLAG, _EXTENDED_FLAG, "first"]
 
 
 def pytest_addoption(parser: pytest.Parser) -> None:

diff --git a/libs/vertexai/tests/integration_tests/test_chat_models.py b/libs/vertexai/tests/integration_tests/test_chat_models.py
@@ -219,6 +219,7 @@ def test_multimodal_media_file_uri(file_uri, mime_type) -> None:
 
 @pytest.mark.release
 @pytest.mark.parametrize("file_uri,mime_type", multimodal_inputs)
+@pytest.mark.first
 def test_multimodal_media_inline_base64(file_uri, mime_type) -> None:
     llm = ChatVertexAI(model_name="gemini-1.5-pro-001", rate_limiter=rate_limiter)
     storage_client = storage.Client()
@@ -239,6 +240,7 @@ def test_multimodal_media_inline_base64(file_uri, mime_type) -> None:
 
 
 @pytest.mark.release
+@pytest.mark.first
 def test_multimodal_media_inline_base64_template() -> None:
     llm = ChatVertexAI(model_name="gemini-1.5-pro-001")
     prompt_template = ChatPromptTemplate.from_messages(
@@ -844,6 +846,7 @@ def test_structured_output_schema_enum():
 
 
 @pytest.mark.extended
+@pytest.mark.first
 def test_context_catching():
     system_instruction = """
     
@@ -903,6 +906,7 @@ def test_context_catching():
 
 
 @pytest.mark.extended
+@pytest.mark.first
 def test_context_catching_tools():
     from langchain import agents
 

diff --git a/libs/vertexai/tests/integration_tests/test_model_garden.py b/libs/vertexai/tests/integration_tests/test_model_garden.py
@@ -81,6 +81,7 @@ def test_model_garden_generate(
 
 @pytest.mark.extended
 @pytest.mark.asyncio
+@pytest.mark.first
 @pytest.mark.parametrize(
     "endpoint_os_variable_name,result_arg",
     [("FALCON_ENDPOINT_ID", "generated_text"), ("LLAMA_ENDPOINT_ID", None)],

diff --git a/libs/vertexai/tests/integration_tests/test_standard.py b/libs/vertexai/tests/integration_tests/test_standard.py
@@ -12,6 +12,7 @@
 rate_limiter = InMemoryRateLimiter(requests_per_second=0.5)
 
 
+@pytest.mark.first
 class TestGeminiAIStandard(ChatModelIntegrationTests):
     @property
     def chat_model_class(self) -> Type[BaseChatModel]: