From 5bb9b41c0e8f9a0072a6052742ba1147766cd35f Mon Sep 17 00:00:00 2001 From: Pratik Khulge Date: Mon, 29 Jun 2026 15:48:29 +0530 Subject: [PATCH] fix(moonshot): apply ruff formatting and fix async tests for CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Run ruff format on batches/main.py and moonshot/batches/handler.py - Replace asyncio.get_event_loop().run_until_complete() with @pytest.mark.asyncio in test_handler.py — the deprecated event-loop getter raises RuntimeError on Python 3.10+ when no loop exists in the thread, causing CI failures --- litellm/batches/main.py | 12 +++-- litellm/llms/moonshot/batches/handler.py | 46 +++++-------------- .../llms/moonshot/batches/test_handler.py | 40 +++++++++------- 3 files changed, 45 insertions(+), 53 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 45de410494b..e3b94b90317 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -351,7 +351,9 @@ def create_batch( @client async def aretrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot" + ] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -397,7 +399,9 @@ def _handle_retrieve_batch_providers_without_provider_config( litellm_params: dict, _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot" + ] = "openai", logging_obj: Optional[Any] = None, ): api_base: Optional[str] = None @@ -532,7 +536,9 @@ def _handle_retrieve_batch_providers_without_provider_config( @client def retrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot" + ] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, diff --git a/litellm/llms/moonshot/batches/handler.py b/litellm/llms/moonshot/batches/handler.py index 0390a51da48..e6303e83bdc 100644 --- a/litellm/llms/moonshot/batches/handler.py +++ b/litellm/llms/moonshot/batches/handler.py @@ -46,11 +46,7 @@ class MoonshotBatchesAPI(BaseLLM): return client resolved_key = api_key or get_secret_str("MOONSHOT_API_KEY") - resolved_base = ( - api_base - or get_secret_str("MOONSHOT_API_BASE") - or MOONSHOT_DEFAULT_API_BASE - ) + resolved_base = api_base or get_secret_str("MOONSHOT_API_BASE") or MOONSHOT_DEFAULT_API_BASE kwargs: dict = {"base_url": resolved_base} if resolved_key is not None: @@ -96,18 +92,13 @@ class MoonshotBatchesAPI(BaseLLM): ) if moonshot_client is None: raise ValueError( - "Moonshot client could not be initialised. " - "Pass api_key or set MOONSHOT_API_KEY in the environment." + "Moonshot client could not be initialised. Pass api_key or set MOONSHOT_API_KEY in the environment." ) if _is_async: if not isinstance(moonshot_client, AsyncOpenAI): - raise ValueError( - "Expected an AsyncOpenAI client for async create_batch." - ) - return self.acreate_batch( - create_batch_data=create_batch_data, client=moonshot_client - ) + raise ValueError("Expected an AsyncOpenAI client for async create_batch.") + return self.acreate_batch(create_batch_data=create_batch_data, client=moonshot_client) response = cast(OpenAI, moonshot_client).batches.create(**create_batch_data) # type: ignore[arg-type] return LiteLLMBatch(**response.model_dump()) @@ -142,18 +133,13 @@ class MoonshotBatchesAPI(BaseLLM): ) if moonshot_client is None: raise ValueError( - "Moonshot client could not be initialised. " - "Pass api_key or set MOONSHOT_API_KEY in the environment." + "Moonshot client could not be initialised. Pass api_key or set MOONSHOT_API_KEY in the environment." ) if _is_async: if not isinstance(moonshot_client, AsyncOpenAI): - raise ValueError( - "Expected an AsyncOpenAI client for async retrieve_batch." - ) - return self.aretrieve_batch( - retrieve_batch_data=retrieve_batch_data, client=moonshot_client - ) + raise ValueError("Expected an AsyncOpenAI client for async retrieve_batch.") + return self.aretrieve_batch(retrieve_batch_data=retrieve_batch_data, client=moonshot_client) response = cast(OpenAI, moonshot_client).batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type] return LiteLLMBatch(**response.model_dump()) @@ -188,18 +174,13 @@ class MoonshotBatchesAPI(BaseLLM): ) if moonshot_client is None: raise ValueError( - "Moonshot client could not be initialised. " - "Pass api_key or set MOONSHOT_API_KEY in the environment." + "Moonshot client could not be initialised. Pass api_key or set MOONSHOT_API_KEY in the environment." ) if _is_async: if not isinstance(moonshot_client, AsyncOpenAI): - raise ValueError( - "Expected an AsyncOpenAI client for async cancel_batch." - ) - return self.acancel_batch( - cancel_batch_data=cancel_batch_data, client=moonshot_client - ) + raise ValueError("Expected an AsyncOpenAI client for async cancel_batch.") + return self.acancel_batch(cancel_batch_data=cancel_batch_data, client=moonshot_client) response = cast(OpenAI, moonshot_client).batches.cancel(**cancel_batch_data) return LiteLLMBatch(**response.model_dump()) @@ -236,15 +217,12 @@ class MoonshotBatchesAPI(BaseLLM): ) if moonshot_client is None: raise ValueError( - "Moonshot client could not be initialised. " - "Pass api_key or set MOONSHOT_API_KEY in the environment." + "Moonshot client could not be initialised. Pass api_key or set MOONSHOT_API_KEY in the environment." ) if _is_async: if not isinstance(moonshot_client, AsyncOpenAI): - raise ValueError( - "Expected an AsyncOpenAI client for async list_batches." - ) + raise ValueError("Expected an AsyncOpenAI client for async list_batches.") return self.alist_batches(client=moonshot_client, after=after, limit=limit) return cast(OpenAI, moonshot_client).batches.list(after=after, limit=limit) # type: ignore diff --git a/tests/test_litellm/llms/moonshot/batches/test_handler.py b/tests/test_litellm/llms/moonshot/batches/test_handler.py index 74dadb945fd..731856f5466 100644 --- a/tests/test_litellm/llms/moonshot/batches/test_handler.py +++ b/tests/test_litellm/llms/moonshot/batches/test_handler.py @@ -11,7 +11,6 @@ real. from __future__ import annotations -import asyncio import os import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -66,7 +65,12 @@ def _sync_client() -> MagicMock: def _async_client() -> MagicMock: - return MagicMock(spec=AsyncOpenAI) + client = MagicMock(spec=AsyncOpenAI) + client.batches.create = AsyncMock() + client.batches.retrieve = AsyncMock() + client.batches.cancel = AsyncMock() + client.batches.list = AsyncMock() + return client # ====================================================================== create @@ -97,10 +101,11 @@ class TestCreateBatch: assert isinstance(result, LiteLLMBatch) assert result.id == "batch-123" - def test_async_create_batch(self): + @pytest.mark.asyncio + async def test_async_create_batch(self): handler = MoonshotBatchesAPI() sdk = _async_client() - sdk.batches.create = AsyncMock(return_value=_sdk_response(_batch_dict())) + sdk.batches.create.return_value = _sdk_response(_batch_dict()) with patch(GET_CLIENT, return_value=sdk): coro = handler.create_batch( @@ -108,7 +113,7 @@ class TestCreateBatch: create_batch_data=CREATE_DATA, **AUTH_KW, ) - result = asyncio.get_event_loop().run_until_complete(coro) + result = await coro sdk.batches.create.assert_called_once_with(**CREATE_DATA) assert isinstance(result, LiteLLMBatch) @@ -143,10 +148,11 @@ class TestRetrieveBatch: sdk.batches.retrieve.assert_called_once_with(**RETRIEVE_DATA) assert isinstance(result, LiteLLMBatch) - def test_async_retrieve_batch(self): + @pytest.mark.asyncio + async def test_async_retrieve_batch(self): handler = MoonshotBatchesAPI() sdk = _async_client() - sdk.batches.retrieve = AsyncMock(return_value=_sdk_response(_batch_dict())) + sdk.batches.retrieve.return_value = _sdk_response(_batch_dict()) with patch(GET_CLIENT, return_value=sdk): coro = handler.retrieve_batch( @@ -154,7 +160,7 @@ class TestRetrieveBatch: retrieve_batch_data=RETRIEVE_DATA, **AUTH_KW, ) - result = asyncio.get_event_loop().run_until_complete(coro) + result = await coro sdk.batches.retrieve.assert_called_once_with(**RETRIEVE_DATA) assert isinstance(result, LiteLLMBatch) @@ -189,10 +195,11 @@ class TestCancelBatch: sdk.batches.cancel.assert_called_once_with(**CANCEL_DATA) assert isinstance(result, LiteLLMBatch) - def test_async_cancel_batch(self): + @pytest.mark.asyncio + async def test_async_cancel_batch(self): handler = MoonshotBatchesAPI() sdk = _async_client() - sdk.batches.cancel = AsyncMock(return_value=_sdk_response(_batch_dict(status="cancelling"))) + sdk.batches.cancel.return_value = _sdk_response(_batch_dict(status="cancelling")) with patch(GET_CLIENT, return_value=sdk): coro = handler.cancel_batch( @@ -200,7 +207,7 @@ class TestCancelBatch: cancel_batch_data=CANCEL_DATA, **AUTH_KW, ) - result = asyncio.get_event_loop().run_until_complete(coro) + result = await coro sdk.batches.cancel.assert_called_once_with(**CANCEL_DATA) assert isinstance(result, LiteLLMBatch) @@ -237,11 +244,12 @@ class TestListBatches: sdk.batches.list.assert_called_once_with(after=None, limit=10) assert result is list_resp - def test_async_list_batches(self): + @pytest.mark.asyncio + async def test_async_list_batches(self): handler = MoonshotBatchesAPI() sdk = _async_client() list_resp = MagicMock() - sdk.batches.list = AsyncMock(return_value=list_resp) + sdk.batches.list.return_value = list_resp with patch(GET_CLIENT, return_value=sdk): coro = handler.list_batches( @@ -250,7 +258,7 @@ class TestListBatches: limit=5, **AUTH_KW, ) - result = asyncio.get_event_loop().run_until_complete(coro) + result = await coro sdk.batches.list.assert_called_once_with(after="cursor-abc", limit=5) assert result is list_resp @@ -274,7 +282,7 @@ class TestGetClient: handler = MoonshotBatchesAPI() with patch("litellm.llms.moonshot.batches.handler.OpenAI") as mock_cls: mock_cls.return_value = MagicMock(spec=OpenAI) - client = handler._get_client( + handler._get_client( api_key=None, api_base="https://api.moonshot.ai/v1", timeout=30.0, @@ -303,7 +311,7 @@ class TestGetClient: handler = MoonshotBatchesAPI() with patch("litellm.llms.moonshot.batches.handler.AsyncOpenAI") as mock_cls: mock_cls.return_value = MagicMock(spec=AsyncOpenAI) - client = handler._get_client( + handler._get_client( api_key="sk-test", api_base="https://api.moonshot.ai/v1", timeout=30.0,