From 5cd8ca236506b954f97af872ec9075df165aa962 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 18:39:32 -0700 Subject: [PATCH] refactor: refactor testing --- tests/batches_tests/test_fine_tuning_api.py | 72 --------- tests/llm_translation/test_rerank.py | 42 ----- .../azure_cancel_expected_output.json | 20 +++ .../azure_cancel_raw_response.json | 18 +++ .../azure_cancel_request.json | 3 + .../azure_create_expected_output.json | 20 +++ .../azure_create_raw_response.json | 18 +++ .../azure_create_request.json | 8 + .../azure_list_raw_response.json | 20 +++ .../azure_list_request.json | 4 + .../llms/azure/test_azure_fine_tuning_api.py | 150 ++++++++++++++++++ 11 files changed, 261 insertions(+), 114 deletions(-) create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_create_request.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_list_request.json create mode 100644 tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py diff --git a/tests/batches_tests/test_fine_tuning_api.py b/tests/batches_tests/test_fine_tuning_api.py index 7e220612ac..034f4e1b49 100644 --- a/tests/batches_tests/test_fine_tuning_api.py +++ b/tests/batches_tests/test_fine_tuning_api.py @@ -123,78 +123,6 @@ async def test_create_fine_tune_jobs_async(): pass -@pytest.mark.asyncio -async def test_azure_create_fine_tune_jobs_async(): - try: - verbose_logger.setLevel(logging.DEBUG) - file_name = "azure_fine_tune.jsonl" - _current_dir = os.path.dirname(os.path.abspath(__file__)) - file_path = os.path.join(_current_dir, file_name) - - file_id = "file-5e4b20ecbd724182b9964f3cd2ab7212" - - create_fine_tuning_response = await litellm.acreate_fine_tuning_job( - model="gpt-35-turbo-1106", - training_file=file_id, - custom_llm_provider="azure", - api_base="https://exampleopenaiendpoint-production.up.railway.app", - ) - - print( - "response from litellm.create_fine_tuning_job=", create_fine_tuning_response - ) - - assert create_fine_tuning_response.id is not None - - # response from Example/mocked endpoint - assert create_fine_tuning_response.model == "davinci-002" - - # list fine tuning jobs - print("listing ft jobs") - ft_jobs = await litellm.alist_fine_tuning_jobs( - limit=2, - custom_llm_provider="azure", - api_base="https://exampleopenaiendpoint-production.up.railway.app", - ) - print("response from litellm.list_fine_tuning_jobs=", ft_jobs) - - # cancel ft job - response = await litellm.acancel_fine_tuning_job( - fine_tuning_job_id=create_fine_tuning_response.id, - custom_llm_provider="azure", - api_key=os.getenv("AZURE_SWEDEN_API_KEY"), - api_base="https://exampleopenaiendpoint-production.up.railway.app", - ) - - print("response from litellm.cancel_fine_tuning_job=", response) - - assert response.status == "cancelled" - assert response.id == create_fine_tuning_response.id - except openai.RateLimitError: - pass - except Exception as e: - if "Job has already completed" in str(e): - pass - else: - pytest.fail(f"Error occurred: {e}") - pass - - -def test_azure_trainingtype_defaults_to_one(): - """ - Azure requires trainingType in extra_body. When omitted, AzureOpenAIFineTuningAPI defaults it to 1. - """ - from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI - - handler = AzureOpenAIFineTuningAPI() - create_data = {"model": "gpt-4o-mini", "training_file": "file-test"} - - handler._ensure_training_type(create_data) - - assert "extra_body" in create_data - assert create_data["extra_body"]["trainingType"] == 1 - - @pytest.mark.asyncio() async def test_create_vertex_fine_tune_jobs_mocked(): load_vertex_ai_credentials() diff --git a/tests/llm_translation/test_rerank.py b/tests/llm_translation/test_rerank.py index edd5981c16..d784677060 100644 --- a/tests/llm_translation/test_rerank.py +++ b/tests/llm_translation/test_rerank.py @@ -148,48 +148,6 @@ async def test_basic_rerank_together_ai(sync_mode): raise e -@pytest.mark.asyncio() -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.skip(reason="Skipping test due to Cohere RBAC issues") -async def test_basic_rerank_azure_ai(sync_mode): - import os - - litellm.set_verbose = True - - if sync_mode is True: - response = litellm.rerank( - model="azure_ai/Cohere-rerank-v3-multilingual-ko", - query="hello", - documents=["hello", "world"], - top_n=3, - api_key=os.getenv("AZURE_AI_COHERE_API_KEY"), - api_base=os.getenv("AZURE_AI_COHERE_API_BASE"), - ) - - print("re rank response: ", response) - - assert response.id is not None - assert response.results is not None - - assert_response_shape(response, custom_llm_provider="together_ai") - else: - response = await litellm.arerank( - model="azure_ai/Cohere-rerank-v3-multilingual-ko", - query="hello", - documents=["hello", "world"], - top_n=3, - api_key=os.getenv("AZURE_AI_COHERE_API_KEY"), - api_base=os.getenv("AZURE_AI_COHERE_API_BASE"), - ) - - print("async re rank response: ", response) - - assert response.id is not None - assert response.results is not None - - assert_response_shape(response, custom_llm_provider="together_ai") - - @pytest.mark.asyncio() @pytest.mark.parametrize("version", ["v1", "v2"]) async def test_rerank_custom_api_base(version): diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json new file mode 100644 index 0000000000..d0a519fcd1 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json @@ -0,0 +1,20 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "cancelled", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3, + "batch_size": null, + "learning_rate_multiplier": null + }, + "organization_id": "", + "result_files": [], + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json new file mode 100644 index 0000000000..0093a04f70 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json @@ -0,0 +1,18 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "canceled", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3 + }, + "organization_id": null, + "result_files": null, + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json new file mode 100644 index 0000000000..bdbbdbad07 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json @@ -0,0 +1,3 @@ +{ + "fine_tuning_job_id": "ftjob-azure-create-123" +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json new file mode 100644 index 0000000000..201ae3047d --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json @@ -0,0 +1,20 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3, + "batch_size": null, + "learning_rate_multiplier": null + }, + "organization_id": "", + "result_files": [], + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json new file mode 100644 index 0000000000..857b249938 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json @@ -0,0 +1,18 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3 + }, + "organization_id": null, + "result_files": null, + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json new file mode 100644 index 0000000000..57319b2698 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json @@ -0,0 +1,8 @@ +{ + "model": "gpt-35-turbo-1106", + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": {}, + "extra_body": { + "trainingType": 1 + } +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json new file mode 100644 index 0000000000..9986e7d4c5 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json @@ -0,0 +1,20 @@ +{ + "object": "list", + "data": [ + { + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running" + }, + { + "id": "ftjob-azure-prev-000", + "object": "fine_tuning.job", + "created_at": 1735603200, + "model": "davinci-002", + "status": "succeeded" + } + ], + "has_more": false +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json new file mode 100644 index 0000000000..6bfbb2ffec --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json @@ -0,0 +1,4 @@ +{ + "after": "ftjob-azure-prev-000", + "limit": 2 +} diff --git a/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py b/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py new file mode 100644 index 0000000000..8d008d6c07 --- /dev/null +++ b/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py @@ -0,0 +1,150 @@ +import json +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest +from openai import AsyncAzureOpenAI + +import litellm +from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI + + +def _expected_dir() -> Path: + return Path(__file__).resolve().parent.parent.parent / "expected_fine_tuning_api" + + +def _load_json(file_name: str) -> dict: + path = _expected_dir() / file_name + assert path.exists(), f"Expected fixture file not found: {path}" + with open(path) as f: + return json.load(f) + + +class _MockSDKResponse: + def __init__(self, payload: dict): + self._payload = payload + + def model_dump(self) -> dict: + return self._payload + + +def _mock_azure_client( + create_payload: dict | None = None, + list_payload: dict | None = None, + cancel_payload: dict | None = None, +): + client = AsyncAzureOpenAI( + api_key="test-key", + api_version="2024-10-21", + azure_endpoint="https://exampleopenaiendpoint-production.up.railway.app", + ) + client.fine_tuning.jobs.create = AsyncMock( + return_value=( + _MockSDKResponse(create_payload) if create_payload is not None else None + ) + ) # type: ignore[method-assign] + client.fine_tuning.jobs.list = AsyncMock( + return_value=list_payload + ) # type: ignore[method-assign] + client.fine_tuning.jobs.cancel = AsyncMock( + return_value=( + _MockSDKResponse(cancel_payload) if cancel_payload is not None else None + ) + ) # type: ignore[method-assign] + return client + + +@pytest.mark.asyncio +async def test_azure_acreate_fine_tuning_job_request_and_output_match_expected_json(): + expected_request = _load_json("azure_create_request.json") + raw_response = _load_json("azure_create_raw_response.json") + expected_output = _load_json("azure_create_expected_output.json") + + mock_client = _mock_azure_client(create_payload=raw_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.acreate_fine_tuning_job( + model="gpt-35-turbo-1106", + training_file="file-5e4b20ecbd724182b9964f3cd2ab7212", + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.create.call_args.kwargs + assert request_kwargs == expected_request + + response_dict = response.model_dump(exclude={"_hidden_params"}) + for key, expected_value in expected_output.items(): + assert key in response_dict, f"Missing key in response: {key}" + assert response_dict[key] == expected_value + + assert response.id is not None + assert response.model == "davinci-002" + + +@pytest.mark.asyncio +async def test_azure_alist_fine_tuning_jobs_request_matches_expected_json(): + expected_request = _load_json("azure_list_request.json") + raw_list_response = _load_json("azure_list_raw_response.json") + + mock_client = _mock_azure_client(list_payload=raw_list_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.alist_fine_tuning_jobs( + after=expected_request["after"], + limit=expected_request["limit"], + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.list.call_args.kwargs + assert request_kwargs == expected_request + assert response == raw_list_response + + +@pytest.mark.asyncio +async def test_azure_acancel_fine_tuning_job_request_and_output_match_expected_json(): + expected_request = _load_json("azure_cancel_request.json") + raw_response = _load_json("azure_cancel_raw_response.json") + expected_output = _load_json("azure_cancel_expected_output.json") + + mock_client = _mock_azure_client(cancel_payload=raw_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.acancel_fine_tuning_job( + fine_tuning_job_id=expected_request["fine_tuning_job_id"], + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.cancel.call_args.kwargs + assert request_kwargs == expected_request + + response_dict = response.model_dump(exclude={"_hidden_params"}) + for key, expected_value in expected_output.items(): + assert key in response_dict, f"Missing key in response: {key}" + assert response_dict[key] == expected_value + + assert response.status == "cancelled" + + +def test_azure_trainingtype_defaults_to_one(): + handler = AzureOpenAIFineTuningAPI() + create_data = {"model": "gpt-4o-mini", "training_file": "file-test"} + + handler._ensure_training_type(create_data) + + assert "extra_body" in create_data + assert create_data["extra_body"]["trainingType"] == 1