litellm/tests/unit/llms/custom_httpx/test_container_handler.py

import json
from importlib.resources import files
from pathlib import Path
from typing import Final
from unittest.mock import MagicMock

import httpx
import pytest
from pydantic import ValidationError

import litellm
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.custom_httpx import container_handler
from litellm.llms.custom_httpx.container_handler import generic_container_handler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.containers.main import DeleteContainerFileResponse
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager

FILE_NOT_FOUND_BODY = {
    "error": {
        "message": "File not found.",
        "type": "invalid_request_error",
        "param": None,
        "code": None,
    }
}


def _sync_client(response: httpx.Response) -> HTTPHandler:
    handler = HTTPHandler()
    handler.client = httpx.Client(transport=httpx.MockTransport(lambda _request: response))
    return handler


def _async_client(response: httpx.Response) -> AsyncHTTPHandler:
    handler = AsyncHTTPHandler()
    handler.client = httpx.AsyncClient(transport=httpx.MockTransport(lambda _request: response))
    return handler


def _handle(endpoint_name: str, client, **overrides):
    return generic_container_handler.handle(
        endpoint_name=endpoint_name,
        container_provider_config=ProviderConfigManager.get_provider_container_config(
            provider=litellm.LlmProviders.OPENAI
        ),
        litellm_params=GenericLiteLLMParams(api_key="sk-test"),
        logging_obj=MagicMock(),
        client=client,
        container_id="cntr_real",
        file_id="cfile_nonexistent",
        **overrides,
    )


def test_binary_endpoint_raises_on_error_status():
    with pytest.raises(BaseLLMException) as exc_info:
        _handle(
            "retrieve_container_file_content",
            _sync_client(httpx.Response(404, json=FILE_NOT_FOUND_BODY)),
        )

    assert exc_info.value.status_code == 404
    assert exc_info.value.message == "File not found."


@pytest.mark.asyncio
async def test_async_binary_endpoint_raises_on_error_status():
    with pytest.raises(BaseLLMException) as exc_info:
        await _handle(
            "aretrieve_container_file_content",
            _async_client(httpx.Response(404, json=FILE_NOT_FOUND_BODY)),
            _is_async=True,
        )

    assert exc_info.value.status_code == 404
    assert exc_info.value.message == "File not found."


def test_binary_endpoint_returns_raw_content_on_success():
    content = _handle(
        "retrieve_container_file_content",
        _sync_client(httpx.Response(200, content=b"\x00binary-payload")),
    )

    assert content == b"\x00binary-payload"


def test_error_status_with_non_json_body_surfaces_response_text():
    with pytest.raises(BaseLLMException) as exc_info:
        _handle(
            "retrieve_container_file_content",
            _sync_client(httpx.Response(502, content=b"<html>bad gateway</html>")),
        )

    assert exc_info.value.status_code == 502
    assert exc_info.value.message == "<html>bad gateway</html>"


def test_json_endpoint_still_raises_provider_error_message():
    with pytest.raises(BaseLLMException) as exc_info:
        _handle(
            "retrieve_container_file",
            _sync_client(httpx.Response(404, json=FILE_NOT_FOUND_BODY)),
        )

    assert exc_info.value.status_code == 404
    assert exc_info.value.message == "File not found."


def test_json_endpoint_sends_the_configured_route_and_parses_its_response_model():
    handler: Final = HTTPHandler()
    handler.client = httpx.Client(
        transport=httpx.MockTransport(
            lambda request: httpx.Response(
                200,
                json={
                    "id": request.url.path,
                    "object": "container.file.deleted",
                    "deleted": request.method == "DELETE",
                },
            )
        )
    )

    response: Final = _handle("delete_container_file", handler)

    assert type(response) is DeleteContainerFileResponse
    assert response.id.endswith("/containers/cntr_real/files/cfile_nonexistent")
    assert response.deleted is True


@pytest.mark.parametrize("missing_key", ["path_params", "query_params"])
def test_an_endpoint_table_without_a_param_list_fails_the_call(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path, missing_key: str
):
    packaged: Final = json.loads(files("litellm.containers").joinpath("endpoints.json").read_text())
    trimmed: Final = {
        "endpoints": [
            {key: value for key, value in entry.items() if key != missing_key} for entry in packaged["endpoints"]
        ]
    }
    (tmp_path / "containers").mkdir()
    (tmp_path / "containers" / "endpoints.json").write_text(json.dumps(trimmed))
    monkeypatch.setattr(container_handler, "__file__", str(tmp_path / "llms" / "custom_httpx" / "container_handler.py"))
    empty_list: Final = {"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}

    with pytest.raises(ValidationError):
        _handle("list_container_files", _sync_client(httpx.Response(200, json=empty_list)))