litellm/tests/unit/containers/test_container_transformation.py

import json
import os
from unittest.mock import MagicMock, patch
import httpx

import pytest


import litellm
from litellm.llms.openai.containers.transformation import OpenAIContainerConfig
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.types.containers.main import (
    ContainerObject,
    ContainerListResponse,
    DeleteContainerResult,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging


class TestOpenAIContainerTransformation:
    """Test suite for OpenAI container transformation functionality."""

    def setup_method(self):
        """Set up test fixtures."""
        self.config = OpenAIContainerConfig()
        self.logging_obj = LiteLLMLogging(
            model="",
            messages=[],
            stream=False,
            call_type="create_container",
            start_time=None,
            litellm_call_id="test_call_id",
            function_id="test_function_id",
        )

    def test_get_supported_openai_params(self):
        """Test that supported OpenAI parameters are returned correctly."""
        supported_params = self.config.get_supported_openai_params()

        # Check that essential container parameters are supported
        assert "name" in supported_params
        assert "expires_after" in supported_params
        assert "file_ids" in supported_params

    def test_map_openai_params_basic(self):
        """Test basic parameter mapping for OpenAI."""
        from litellm.types.containers.main import ContainerCreateOptionalRequestParams

        optional_params = ContainerCreateOptionalRequestParams(
            {
                "expires_after": {"anchor": "last_active_at", "minutes": 30},
                "file_ids": ["file_1", "file_2"],
            }
        )

        mapped_params = self.config.map_openai_params(
            optional_params, drop_params=False
        )

        assert mapped_params["expires_after"]["minutes"] == 30
        assert mapped_params["file_ids"] == ["file_1", "file_2"]

    def test_validate_environment(self):
        """Test environment validation adds proper headers."""
        headers = {}
        api_key = "sk-test123"

        validated_headers = self.config.validate_environment(
            headers=headers, api_key=api_key
        )

        assert "Authorization" in validated_headers
        assert validated_headers["Authorization"] == f"Bearer {api_key}"
        # Note: Content-Type is not added by validate_environment method

    def test_get_complete_url(self):
        """Test complete URL generation."""
        api_base = "https://api.openai.com/v1"
        litellm_params = {}

        url = self.config.get_complete_url(
            api_base=api_base, litellm_params=litellm_params
        )

        assert url == "https://api.openai.com/v1/containers"

    def test_get_complete_url_with_custom_base(self):
        """Test complete URL generation with custom API base."""
        api_base = "https://custom.openai.com/v1"
        litellm_params = {}

        url = self.config.get_complete_url(
            api_base=api_base, litellm_params=litellm_params
        )

        assert url == "https://custom.openai.com/v1/containers"

    def test_transform_container_create_request(self):
        """Test container create request transformation."""
        from litellm.types.router import GenericLiteLLMParams

        litellm_params = GenericLiteLLMParams()
        headers = {"Authorization": "Bearer sk-test123"}
        name = "Test Container"
        container_create_optional_request_params = {
            "expires_after": {"anchor": "last_active_at", "minutes": 20},
            "file_ids": ["file_123"],
        }

        data = self.config.transform_container_create_request(
            name=name,
            container_create_optional_request_params=container_create_optional_request_params,
            litellm_params=litellm_params,
            headers=headers,
        )

        assert data["name"] == name
        assert (
            data["expires_after"]
            == container_create_optional_request_params["expires_after"]
        )
        assert data["file_ids"] == container_create_optional_request_params["file_ids"]

    def test_transform_container_create_response(self):
        """Test container create response transformation."""
        # Mock HTTP response
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = {
            "id": "cntr_123456",
            "object": "container",
            "created_at": 1747857508,
            "status": "running",
            "expires_after": {"anchor": "last_active_at", "minutes": 20},
            "last_active_at": 1747857508,
            "name": "Test Container",
        }

        container = self.config.transform_container_create_response(
            raw_response=mock_response, logging_obj=self.logging_obj
        )

        assert isinstance(container, ContainerObject)
        assert container.id == "cntr_123456"
        assert container.name == "Test Container"
        assert container.status == "running"
        assert container.object == "container"

    def test_transform_container_list_request(self):
        """Test container list request transformation."""
        api_base = "https://api.openai.com/v1/containers"
        litellm_params = {}
        headers = {"Authorization": "Bearer sk-test123"}
        after = "cntr_123"
        limit = 10
        order = "desc"

        url, params = self.config.transform_container_list_request(
            api_base=api_base,
            litellm_params=litellm_params,
            headers=headers,
            after=after,
            limit=limit,
            order=order,
        )

        assert url == api_base
        assert params["after"] == after
        assert params["limit"] == str(limit)  # Should be string for query params
        assert params["order"] == order

    def test_transform_container_list_response(self):
        """Test container list response transformation."""
        # Mock HTTP response
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = {
            "object": "list",
            "data": [
                {
                    "id": "cntr_1",
                    "object": "container",
                    "created_at": 1747857508,
                    "status": "running",
                    "expires_after": {"anchor": "last_active_at", "minutes": 20},
                    "last_active_at": 1747857508,
                    "name": "Container 1",
                },
                {
                    "id": "cntr_2",
                    "object": "container",
                    "created_at": 1747857600,
                    "status": "running",
                    "expires_after": {"anchor": "last_active_at", "minutes": 15},
                    "last_active_at": 1747857600,
                    "name": "Container 2",
                },
            ],
            "first_id": "cntr_1",
            "last_id": "cntr_2",
            "has_more": False,
        }

        container_list = self.config.transform_container_list_response(
            raw_response=mock_response, logging_obj=self.logging_obj
        )

        assert isinstance(container_list, ContainerListResponse)
        assert len(container_list.data) == 2
        assert container_list.first_id == "cntr_1"
        assert container_list.last_id == "cntr_2"
        assert container_list.has_more == False

    def test_transform_container_retrieve_request(self):
        """Test container retrieve request transformation."""
        container_id = "cntr_test123"
        api_base = "https://api.openai.com/v1/containers"
        litellm_params = {}
        headers = {"Authorization": "Bearer sk-test123"}

        url, params = self.config.transform_container_retrieve_request(
            container_id=container_id,
            api_base=api_base,
            litellm_params=litellm_params,
            headers=headers,
        )

        assert url == f"{api_base}/{container_id}"
        assert params == {}  # No query params for retrieve

    def test_transform_container_retrieve_request_encodes_path_traversal(self):
        """Test container IDs are treated as a single upstream path segment."""
        api_base = "https://api.openai.com/v1/containers"

        url, params = self.config.transform_container_retrieve_request(
            container_id="../../vector_stores?x=1#frag",
            api_base=api_base,
            litellm_params={},
            headers={},
        )

        assert (
            url
            == "https://api.openai.com/v1/containers/..%2F..%2Fvector_stores%3Fx%3D1%23frag"
        )
        assert params == {}

    def test_transform_container_retrieve_response(self):
        """Test container retrieve response transformation."""
        # Mock HTTP response
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = {
            "id": "cntr_retrieve_123",
            "object": "container",
            "created_at": 1747857508,
            "status": "running",
            "expires_after": {"anchor": "last_active_at", "minutes": 20},
            "last_active_at": 1747857508,
            "name": "Retrieved Container",
        }

        container = self.config.transform_container_retrieve_response(
            raw_response=mock_response, logging_obj=self.logging_obj
        )

        assert isinstance(container, ContainerObject)
        assert container.id == "cntr_retrieve_123"
        assert container.name == "Retrieved Container"

    def test_transform_container_delete_request(self):
        """Test container delete request transformation."""
        container_id = "cntr_delete_123"
        api_base = "https://api.openai.com/v1/containers"
        litellm_params = {}
        headers = {"Authorization": "Bearer sk-test123"}

        url, params = self.config.transform_container_delete_request(
            container_id=container_id,
            api_base=api_base,
            litellm_params=litellm_params,
            headers=headers,
        )

        assert url == f"{api_base}/{container_id}"
        assert params == {}  # No query params for delete

    def test_transform_container_delete_response(self):
        """Test container delete response transformation."""
        # Mock HTTP response
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = {
            "id": "cntr_delete_123",
            "object": "container.deleted",
            "deleted": True,
        }

        delete_result = self.config.transform_container_delete_response(
            raw_response=mock_response, logging_obj=self.logging_obj
        )

        assert isinstance(delete_result, DeleteContainerResult)
        assert delete_result.id == "cntr_delete_123"
        assert delete_result.object == "container.deleted"
        assert delete_result.deleted == True

    def test_get_error_class(self):
        """Test error class handling."""
        import httpx
        from litellm.llms.base_llm.chat.transformation import BaseLLMException

        with pytest.raises(BaseLLMException) as exc_info:
            self.config.get_error_class(
                error_message="Test error", status_code=400, headers={}
            )

        assert "Test error" in str(exc_info.value)

    def test_transform_with_none_optional_params(self):
        """Test transformation handles None optional parameters correctly."""
        from litellm.types.router import GenericLiteLLMParams

        litellm_params = GenericLiteLLMParams()
        headers = {"Authorization": "Bearer sk-test123"}
        name = "Test Container"
        container_create_optional_request_params = {
            "expires_after": None,
            "file_ids": None,
        }

        data = self.config.transform_container_create_request(
            name=name,
            container_create_optional_request_params=container_create_optional_request_params,
            litellm_params=litellm_params,
            headers=headers,
        )

        assert data["name"] == name
        # None values should be included as None
        assert data["expires_after"] is None
        assert data["file_ids"] is None

    def test_container_create_response_includes_cost(self, monkeypatch):
        """Test that container create response includes code interpreter cost calculation."""
        # Force use of local model cost map for CI/CD consistency
        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
        litellm.model_cost = litellm.get_model_cost_map(url="")

        from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
            StandardBuiltInToolCostTracking,
        )

        # Mock HTTP response
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = {
            "id": "cntr_cost_test",
            "object": "container",
            "created_at": 1747857508,
            "status": "running",
            "expires_after": {"anchor": "last_active_at", "minutes": 20},
            "last_active_at": 1747857508,
            "name": "Cost Test Container",
        }

        # Transform the response
        container = self.config.transform_container_create_response(
            raw_response=mock_response, logging_obj=self.logging_obj
        )

        # Verify the container object is created
        assert isinstance(container, ContainerObject)
        assert container.id == "cntr_cost_test"

        # Verify that _hidden_params contains cost information
        assert hasattr(container, "_hidden_params")
        assert container._hidden_params is not None
        assert "additional_headers" in container._hidden_params
        assert (
            "llm_provider-x-litellm-response-cost"
            in container._hidden_params["additional_headers"]
        )

        expected_cost = StandardBuiltInToolCostTracking.get_cost_for_code_interpreter(
            sessions=1, provider="openai"
        )
        actual_cost = container._hidden_params["additional_headers"][
            "llm_provider-x-litellm-response-cost"
        ]

        assert actual_cost == expected_cost