From 6b8153fc49e1b31d8561f8d9f43e48ad5371cbc6 Mon Sep 17 00:00:00 2001 From: Ratish Panda Date: Tue, 1 Apr 2025 03:22:57 +0530 Subject: [PATCH 1/3] add test for chatgooglegenerativeai transport check --- .../integration_tests/test_chat_models.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/libs/genai/tests/integration_tests/test_chat_models.py b/libs/genai/tests/integration_tests/test_chat_models.py index efc1a773..5595452e 100644 --- a/libs/genai/tests/integration_tests/test_chat_models.py +++ b/libs/genai/tests/integration_tests/test_chat_models.py @@ -574,3 +574,21 @@ async def model_astream(context: str) -> List[BaseMessageChunk]: result = asyncio.run(model_astream("How can you help me?")) assert len(result) > 0 assert isinstance(result[0], AIMessageChunk) + +@pytest.mark.extended +def test_prediction_client_transport(): + model = ChatGoogleGenerativeAI(model=_MODEL) + assert model.client.transport.kind == "grpc" + + model = ChatGoogleGenerativeAI(model=_MODEL, transport="rest") + assert model.client.transport.kind == "rest" + + async def check_async_client(): + model = ChatGoogleGenerativeAI(model=_MODEL) + assert model.async_client.transport.kind == "grpc_asyncio" + + # test auto conversion of transport to "grpc_asyncio" from "rest" + model = ChatGoogleGenerativeAI(model=_MODEL, transport="rest") + assert model.async_client.transport.kind == "grpc_asyncio" + + asyncio.run(check_async_client()) \ No newline at end of file From 91611a4c9d7aa000c4448de99067d0951a40aad6 Mon Sep 17 00:00:00 2001 From: Ratish Panda Date: Tue, 1 Apr 2025 17:05:02 +0530 Subject: [PATCH 2/3] linting/format --- libs/genai/tests/integration_tests/test_chat_models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/genai/tests/integration_tests/test_chat_models.py b/libs/genai/tests/integration_tests/test_chat_models.py index 5595452e..8daf541b 100644 --- a/libs/genai/tests/integration_tests/test_chat_models.py +++ b/libs/genai/tests/integration_tests/test_chat_models.py @@ -575,6 +575,7 @@ async def model_astream(context: str) -> List[BaseMessageChunk]: assert len(result) > 0 assert isinstance(result[0], AIMessageChunk) + @pytest.mark.extended def test_prediction_client_transport(): model = ChatGoogleGenerativeAI(model=_MODEL) @@ -591,4 +592,4 @@ async def check_async_client(): model = ChatGoogleGenerativeAI(model=_MODEL, transport="rest") assert model.async_client.transport.kind == "grpc_asyncio" - asyncio.run(check_async_client()) \ No newline at end of file + asyncio.run(check_async_client()) From 44a2454bd863971e40f652dc29d6924e77c68048 Mon Sep 17 00:00:00 2001 From: Ratish Panda Date: Thu, 3 Apr 2025 01:25:58 +0530 Subject: [PATCH 3/3] another lint issue fix --- libs/genai/tests/integration_tests/test_chat_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/genai/tests/integration_tests/test_chat_models.py b/libs/genai/tests/integration_tests/test_chat_models.py index 8daf541b..b79e8a64 100644 --- a/libs/genai/tests/integration_tests/test_chat_models.py +++ b/libs/genai/tests/integration_tests/test_chat_models.py @@ -577,14 +577,14 @@ async def model_astream(context: str) -> List[BaseMessageChunk]: @pytest.mark.extended -def test_prediction_client_transport(): +def test_prediction_client_transport() -> None: model = ChatGoogleGenerativeAI(model=_MODEL) assert model.client.transport.kind == "grpc" model = ChatGoogleGenerativeAI(model=_MODEL, transport="rest") assert model.client.transport.kind == "rest" - async def check_async_client(): + async def check_async_client() -> None: model = ChatGoogleGenerativeAI(model=_MODEL) assert model.async_client.transport.kind == "grpc_asyncio"