scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

new mocked

Browse source
This commit is contained in:
naaa760 2025-11-25 17:28:02 +05:30
parent a24b43cbdf
commit 2cf86e8ef8
1 changed files with 44 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
tests/batches_tests/test_openai_batches_and_files.py
View file

@ -447,6 +447,18 @@ mock_vertex_batch_response = {
"completionStats": {"successfulCount": 0, "failedCount": 0, "remainingCount": 100},
}
mock_vertex_list_response = {
"batchPredictionJobs": [
mock_vertex_batch_response,
{
**mock_vertex_batch_response,
"name": "projects/123456789/locations/us-central1/batchPredictionJobs/test-batch-id-789",
"state": "JOB_STATE_SUCCEEDED",
},
],
"nextPageToken": "",
}
@pytest.mark.asyncio
async def test_avertex_batch_prediction(monkeypatch):
@ -533,3 +545,35 @@ async def test_avertex_batch_prediction(monkeypatch):
print("retrieved_batch=", retrieved_batch)
assert retrieved_batch.id == "test-batch-id-456"
@pytest.mark.asyncio
async def test_vertex_list_batches(monkeypatch):
monkeypatch.setenv("GCS_BUCKET_NAME", "litellm-local")
monkeypatch.setenv("VERTEXAI_PROJECT", "litellm-test-project")
monkeypatch.setenv("VERTEXAI_LOCATION", "us-central1")
monkeypatch.setattr(
"litellm.llms.vertex_ai.batches.handler.VertexAIBatchPrediction._ensure_access_token",
lambda self, credentials, project_id, custom_llm_provider: ("mock-token", "litellm-test-project"),
)
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
) as mock_get:
mock_get_response = MagicMock()
mock_get_response.json.return_value = mock_vertex_list_response
mock_get_response.status_code = 200
mock_get_response.raise_for_status.return_value = None
mock_get.return_value = mock_get_response
list_response = await litellm.alist_batches(
custom_llm_provider="vertex_ai",
limit=2,
)
assert list_response["object"] == "list"
assert list_response["has_more"] is False
assert len(list_response["data"]) == 2
assert list_response["data"][0].id == "test-batch-id-456"
assert list_response["data"][1].id == "test-batch-id-789"