Skip to content

Commit 152daeb

Browse files
imperorrpmdrxy
andauthored
test(genai): add test for ChatGoogleGenerativeAI class transport check (#831)
<!-- # Thank you for contributing to LangChain-google! --> <!-- ## Checklist for PR Creation - [ ] PR Title: "[package]: [brief description]" - Where "package" is genai, vertexai, or community - Use "docs: ..." for purely docs changes, "templates: ..." for template changes, "infra: ..." for CI changes - Example: "community: add foobar LLM" - [ ] PR Description and Relevant issues: - Description of the change - Relevant issues (if applicable) - Any dependencies required for this change - [ ] Add Tests and Docs: - If adding a new integration: 1. Include a test for the integration (preferably unit tests that do not rely on network access) 2. Add an example notebook showing its use (place in the `docs/docs/integrations` directory) - [ ] Lint and Test: - Run `make format`, `make lint`, and `make test` from the root of the package(s) you've modified - See contribution guidelines for more: https://github.com/langchain-ai/langchain-google/blob/main/README.md#contribute-code --> <!-- ## Additional guidelines - [ ] PR title and description are appropriate - [ ] Necessary tests and documentation have been added - [ ] Lint and tests pass successfully - [ ] The following additional guidelines are adhered to: - Optional dependencies are imported within functions - No unnecessary dependencies added to pyproject.toml files (except those required for unit tests) - PR doesn't touch more than one package - Changes are backwards compatible --> ## PR Description <!-- e.g. "Implement user authentication feature" --> Added a test for PR #822 for parity with the VertexAI transport check test at https://github.com/imperorrp/langchain-google/blob/main/libs/vertexai/tests/integration_tests/test_chat_models.py#L788-L803 ## Relevant issues <!-- e.g. "Fixes #000" --> #791 ## Type <!-- Select the type of Pull Request --> <!-- Keep only the necessary ones --> ✅ Test ## Testing <!-- Test procedure --> <!-- Test result --> ![image](https://github.com/user-attachments/assets/0b4304b7-1b7c-48e3-9710-fac476500911) --------- Co-authored-by: Mason Daugherty <[email protected]>
1 parent 7ed36af commit 152daeb

File tree

1 file changed

+19
-0
lines changed

1 file changed

+19
-0
lines changed

libs/genai/tests/integration_tests/test_chat_models.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -833,6 +833,25 @@ async def model_astream(context: str) -> list[BaseMessageChunk]:
833833
assert isinstance(result[0], AIMessageChunk)
834834

835835

836+
@pytest.mark.extended
837+
def test_prediction_client_transport() -> None:
838+
model = ChatGoogleGenerativeAI(model=_MODEL)
839+
assert model.client.transport.kind == "grpc"
840+
841+
model = ChatGoogleGenerativeAI(model=_MODEL, transport="rest")
842+
assert model.client.transport.kind == "rest"
843+
844+
async def check_async_client() -> None:
845+
model = ChatGoogleGenerativeAI(model=_MODEL)
846+
assert model.async_client.transport.kind == "grpc_asyncio"
847+
848+
# test auto conversion of transport to "grpc_asyncio" from "rest"
849+
model = ChatGoogleGenerativeAI(model=_MODEL, transport="rest")
850+
assert model.async_client.transport.kind == "grpc_asyncio"
851+
852+
asyncio.run(check_async_client())
853+
854+
836855
def test_search_builtin() -> None:
837856
llm = ChatGoogleGenerativeAI(model="models/gemini-2.0-flash-001").bind_tools(
838857
[{"google_search": {}}]

0 commit comments

Comments
 (0)