Merge pull request #22763 from BerriAI/litellm_test_e2e_batches_test

feat(tests): add proxy e2e azure batches test

.circleci/config.yml

@@ -3689,6 +3689,114 @@ jobs:
       - store_test_results:
           path: test-results
 
+  proxy_e2e_azure_batches_tests:
+    machine:
+      image: ubuntu-2204:2023.10.1
+    resource_class: xlarge
+    working_directory: ~/project
+    steps:
+      - checkout
+      - setup_google_dns
+      - run:
+          name: Install Docker CLI
+          command: |
+            curl -fsSL https://get.docker.com | sh
+            sudo usermod -aG docker $USER
+            docker version
+      - run:
+          name: Install Python 3.12
+          command: |
+            curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh
+            bash miniconda.sh -b -p $HOME/miniconda
+            export PATH="$HOME/miniconda/bin:$PATH"
+            conda init bash
+            source ~/.bashrc
+            conda create -n myenv python=3.12 -y
+            conda activate myenv
+            python --version
+      - run:
+          name: Install Poetry
+          command: |
+            export PATH="$HOME/miniconda/bin:$PATH"
+            source $HOME/miniconda/etc/profile.d/conda.sh
+            conda activate myenv
+            pip install poetry
+      - run:
+          name: Install dockerize
+          command: |
+            wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz
+            sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz
+            rm dockerize-linux-amd64-v0.6.1.tar.gz
+      - run:
+          name: Start PostgreSQL Database
+          command: |
+            docker run -d \
+              --name postgres-db \
+              -e POSTGRES_USER=llmproxy \
+              -e POSTGRES_PASSWORD=dbpassword9090 \
+              -e POSTGRES_DB=litellm \
+              -p 5432:5432 \
+              postgres:15
+      - run:
+          name: Wait for PostgreSQL to be ready
+          command: dockerize -wait tcp://localhost:5432 -timeout 1m
+      - run:
+          name: Install system dependencies
+          command: |
+            sudo apt-get update -y
+            sudo apt-get install -y libpq-dev
+      - run:
+          name: Install Dependencies
+          command: |
+            export PATH="$HOME/miniconda/bin:$PATH"
+            source $HOME/miniconda/etc/profile.d/conda.sh
+            conda activate myenv
+            poetry config virtualenvs.in-project true
+            poetry install --with dev,proxy-dev --extras "proxy"
+            poetry run pip install psycopg2-binary uvicorn fastapi httpx tenacity
+      - run:
+          name: Setup litellm-enterprise
+          command: |
+            export PATH="$HOME/miniconda/bin:$PATH"
+            source $HOME/miniconda/etc/profile.d/conda.sh
+            conda activate myenv
+            poetry run pip install --force-reinstall --no-deps -e enterprise/
+      - run:
+          name: Generate Prisma client
+          command: |
+            export PATH="$HOME/miniconda/bin:$PATH"
+            source $HOME/miniconda/etc/profile.d/conda.sh
+            conda activate myenv
+            poetry run prisma generate --schema litellm/proxy/schema.prisma
+      - run:
+          name: Run Prisma migrations
+          command: |
+            export PATH="$HOME/miniconda/bin:$PATH"
+            source $HOME/miniconda/etc/profile.d/conda.sh
+            conda activate myenv
+            export DATABASE_URL=postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
+            cd litellm/proxy
+            poetry run prisma migrate deploy --schema schema.prisma
+            cd ../..
+      - run:
+          name: Run Azure Batch E2E Tests
+          command: |
+            export PATH="$HOME/miniconda/bin:$PATH"
+            source $HOME/miniconda/etc/profile.d/conda.sh
+            conda activate myenv
+            export DATABASE_URL=postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
+            export USE_LOCAL_LITELLM=true
+            export USE_MOCK_MODELS=true
+            export USE_STATE_TRACKER=true
+            export LITELLM_LOG=DEBUG
+            poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py \
+              -vv -s -k "test_e2e_managed_batch" \
+              --tb=short \
+              --maxfail=3 \
+              --durations=10 \
+              --junitxml=test-results/junit.xml
+          no_output_timeout: 30m
+
   upload-coverage:
     docker:
       - image: cimg/python:3.9
@@ -4458,6 +4566,12 @@ workflows:
               only:
                 - main
                 - /litellm_.*/
+      - proxy_e2e_azure_batches_tests:
+          filters:
+            branches:
+              only:
+                - main
+                - /litellm_.*/
       - llm_translation_testing:
           filters:
             branches:

.github/workflows/test-proxy-e2e-azure-batches.yml

@@ -0,0 +1,90 @@
+name: Proxy E2E Azure Batches Tests
+
+on:
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  proxy_e2e_azure_batches_tests:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+
+    services:
+      postgres:
+        image: postgres:15
+        env:
+          POSTGRES_USER: llmproxy
+          POSTGRES_PASSWORD: dbpassword9090
+          POSTGRES_DB: litellm
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+
+      - name: Cache Poetry dependencies
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cache/pypoetry
+            ~/.cache/pip
+            .venv
+          key: ${{ runner.os }}-poetry-e2e-batches-${{ hashFiles('poetry.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-poetry-e2e-batches-
+            ${{ runner.os }}-poetry-
+
+      - name: Install dependencies
+        run: |
+          poetry config virtualenvs.in-project true
+          poetry install --with dev,proxy-dev --extras "proxy"
+          poetry run pip install psycopg2-binary uvicorn fastapi httpx
+
+      - name: Setup litellm-enterprise
+        run: |
+          poetry run pip install --force-reinstall --no-deps -e enterprise/
+
+      - name: Generate Prisma client
+        run: |
+          poetry run prisma generate --schema litellm/proxy/schema.prisma
+
+      - name: Run Prisma migrations
+        env:
+          DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
+        run: |
+          cd litellm/proxy
+          poetry run prisma migrate deploy --schema schema.prisma
+          cd ../..
+
+      - name: Run Azure Batch E2E Tests
+        env:
+          DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
+          USE_LOCAL_LITELLM: "true"
+          USE_MOCK_MODELS: "true"
+          USE_STATE_TRACKER: "true"
+          LITELLM_LOG: DEBUG
+        run: |
+          poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py \
+            -vv -s -k "test_e2e_managed_batch" \
+            --tb=short \
+            --maxfail=3 \
+            --durations=10
+

tests/proxy_e2e_azure_batches_tests/base_integration_test.py

@@ -0,0 +1,494 @@
+"""Base class for LiteLLM integration tests.
+
+Supports both local (mock) and remote testing modes via environment variables:
+- USE_LOCAL_LITELLM: When "true", uses local LiteLLM at localhost:4000 (default: false)
+- USE_MOCK_MODELS: When "true", uses mock model names (default: false)
+- LITELLM_API_KEY: API key for remote LiteLLM (required when USE_LOCAL_LITELLM=false)
+- LITELLM_BASE_URL: Base URL for remote LiteLLM (required when USE_LOCAL_LITELLM=false)
+"""
+
+import enum
+import os
+import time
+import uuid
+from abc import ABC
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Tuple, Union
+
+import httpx
+import openai
+import pytest
+import requests
+from urllib3.exceptions import InsecureRequestWarning
+
+requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)
+
+LOCAL_LITELLM_BASE_URL = "http://localhost:4000"
+LOCAL_MOCK_SERVER_URL = "http://localhost:8090"
+
+if "USE_LOCAL_LITELLM" not in os.environ:
+    os.environ["USE_LOCAL_LITELLM"] = "true"
+if "USE_MOCK_MODELS" not in os.environ:
+    os.environ["USE_MOCK_MODELS"] = "true"
+if "USE_STATE_TRACKER" not in os.environ:
+    os.environ["USE_STATE_TRACKER"] = "true"
+if "DATABASE_URL" not in os.environ:
+    os.environ["DATABASE_URL"] = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm"
+
+
+def use_local_litellm() -> bool:
+    return os.environ.get("USE_LOCAL_LITELLM", "false").lower() == "true"
+
+
+def use_remote_litellm() -> bool:
+    return not use_local_litellm()
+
+
+def use_mock_models() -> bool:
+    return os.environ.get("USE_MOCK_MODELS", "false").lower() == "true"
+
+
+def get_local_litellm_base_url() -> str:
+    return LOCAL_LITELLM_BASE_URL
+
+
+def get_remote_litellm_base_url() -> str:
+    return os.environ.get("LITELLM_BASE_URL", "").rstrip("/")
+
+
+def get_litellm_base_url() -> str:
+    if use_local_litellm():
+        return get_local_litellm_base_url()
+    return get_remote_litellm_base_url()
+
+
+def get_litellm_api_key() -> str:
+    if use_local_litellm():
+        return "sk-1234"
+    return os.environ.get("LITELLM_API_KEY", "")
+
+
+def get_mock_server_base_url() -> str:
+    return LOCAL_MOCK_SERVER_URL
+
+
+def get_responses_model_name() -> str:
+    if use_mock_models():
+        return "openai-fake-gpt-4o"
+    return "gpt-4o-mini-2024-07-18"
+
+
+def model_id(param) -> str:
+    """Generate a test ID from a model name or tuple containing model name.
+
+    Handles both:
+    - String: "gpt-4o-mini" -> "gpt_4o_mini"
+    - Tuple: ("gpt-4o", "openai/gpt-4o") -> "gpt_4o"
+    """
+    if isinstance(param, tuple):
+        name = param[0]
+    else:
+        name = param
+    return name.replace("-", "_").replace(".", "_")
+
+
+def generate_test_id(
+    params: Tuple[str, ...],
+    test_name: str = "test",
+) -> str:
+    """Generate test ID from model parameters tuple.
+
+    Handles two tuple formats:
+    - 6 elements: (provider, deployment, model_name, api_version, action, reason)
+    - 7 elements: (provider, deployment, model_name, api_version, model_id, action, reason)
+
+    Uses model_id (position 4) if 7 elements, otherwise model_name (position 2).
+    """
+    provider = params[0]
+    deployment = params[1]
+    api_version = params[3]
+
+    if len(params) == 7:
+        identifier = params[4]  # model_id
+    else:
+        identifier = params[2]  # model_name
+
+    test_id = "/".join([provider, deployment, api_version, identifier, test_name])
+    return test_id.replace("-", "_").replace(".", "_")
+
+
+class ModelTestAction(enum.Enum):
+    NOT_APPLICABLE = 1
+    SKIP = 2
+    RUN = 3
+    WARN_ON_FAIL = 4
+
+    def applicable(self) -> bool:
+        return self.value != ModelTestAction.NOT_APPLICABLE.value
+
+
+class BaseLiteLLMIntegrationTest(ABC):
+    """Base class for all LiteLLM integration tests.
+
+    Supports both local/mock and remote testing based on environment variables.
+    """
+
+    @staticmethod
+    def get_api_key() -> str:
+        return get_litellm_api_key()
+
+    @staticmethod
+    def get_base_url() -> str:
+        return get_litellm_base_url()
+
+    @staticmethod
+    def get_ca_bundle_path() -> str:
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        # change if needed
+
+    @classmethod
+    def _get_ssl_verify_setting(cls) -> Union[bool, str]:
+        """Get the appropriate SSL verification setting based on mode.
+
+        Returns path string (not SSLContext) for compatibility with both
+        requests and httpx libraries.
+        """
+        if use_local_litellm():
+            return False
+        ca_bundle_path = cls.get_ca_bundle_path()
+        if os.path.exists(ca_bundle_path):
+            return ca_bundle_path
+        return True
+
+    @classmethod
+    def setup_class(cls):
+        cls.api_key = cls.get_api_key()
+        cls.base_url = cls.get_base_url()
+
+        if not cls.api_key:
+            pytest.fail(
+                "API key is not available. Set LITELLM_API_KEY or USE_LOCAL_LITELLM=true",
+            )
+        if not cls.base_url:
+            pytest.fail(
+                "Base URL is not available. Set LITELLM_BASE_URL or USE_LOCAL_LITELLM=true",
+            )
+
+        verify_setting = cls._get_ssl_verify_setting()
+
+        if use_remote_litellm() and isinstance(verify_setting, str):
+            os.environ["REQUESTS_CA_BUNDLE"] = verify_setting
+            os.environ["CURL_CA_BUNDLE"] = verify_setting
+            print(f"Using CA bundle: {verify_setting}")
+
+        cls.openai_client = openai.OpenAI(
+            base_url=cls.base_url,
+            api_key=cls.api_key,
+            http_client=httpx.Client(verify=verify_setting),
+        )
+
+    @classmethod
+    def make_request(
+        cls,
+        method: str,
+        endpoint: str,
+        timeout_secs: int,
+        **kwargs,
+    ) -> requests.Response:
+        headers = kwargs.get("headers", {})
+        headers["Authorization"] = f"Bearer {cls.api_key}"
+        kwargs["headers"] = headers
+        kwargs.setdefault("timeout", timeout_secs)
+        kwargs.setdefault("verify", cls._get_ssl_verify_setting())
+
+        url = f"{cls.base_url}{endpoint}"
+        return requests.request(method, url, **kwargs)
+
+    @staticmethod
+    def generate_request_id() -> str:
+        return f"req-{uuid.uuid4().hex[:8]}"
+
+    @staticmethod
+    def get_timeout_secs(model_name: str) -> int:
+        model_lower = model_name.lower()
+        slow_models = ["gpt-5", "gpt_5", "o1", "claude-opus", "claude_opus", "o3", "o4"]
+
+        if any(slow_model in model_lower for slow_model in slow_models):
+            return 300
+        return 60
+
+    @staticmethod
+    def generate_unique_filename(extension: str = "txt") -> str:
+        return f"test_{time.time()}.{extension}"
+
+    @staticmethod
+    def extract_model_params(model_data: Dict[str, Any]) -> Tuple[str, str, str, str]:
+        """Extract standardized parameters from model data."""
+        model_name = model_data.get("model_name", "")
+        model_info = model_data.get("model_info", {})
+        provider = model_info.get("litellm_provider", "unknown")
+        litellm_params = model_data.get("litellm_params", {})
+
+        if provider == "azure":
+            api_base = litellm_params.get("api_base", "unknown")
+            if api_base != "unknown" and "//" in api_base:
+                domain_name = api_base.split("//")[1]
+                deployment = domain_name.split(".")[0]
+            else:
+                deployment = "unknown"
+            api_version = litellm_params.get("api_version", "unknown")
+        elif provider in ["bedrock", "bedrock_converse"]:
+            deployment = litellm_params.get("aws_region_name", "unknown")
+            api_version = "unknown"
+        else:
+            deployment = "unknown"
+            api_version = "unknown"
+
+        return provider, deployment, model_name, api_version
+
+    @classmethod
+    def _fetch_all_models_from_litellm(cls) -> List[Dict[str, Any]]:
+        base_url = cls.get_base_url()
+        api_key = cls.get_api_key()
+
+        if not api_key or not base_url:
+            return []
+
+        verify_setting = cls._get_ssl_verify_setting()
+
+        response = requests.get(
+            f"{base_url}/model/info",
+            headers={"Authorization": f"Bearer {api_key}"},
+            verify=verify_setting,
+            timeout=30,
+        )
+
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to fetch all models from {base_url}. Response code: {response.status_code}",
+            )
+
+        data = response.json()
+        return data.get("data", [])
+
+    @classmethod
+    def _fetch_all_approved_models(cls) -> List[Dict[str, Any]]:
+        return cls._fetch_all_models_from_litellm()
+
+    @classmethod
+    def build_model_test_params(
+        cls,
+        should_skip_model: Callable[
+            [str, str, str, str, Dict[str, Any]],
+            Tuple["ModelTestAction", str],
+        ],
+        include_model_id: bool = False,
+        include_load_balanced: bool = False,
+    ) -> List[Tuple[str, ...]]:
+        """Build test parameters from all approved models.
+
+        Args:
+            should_skip_model: Callback that determines if a model should be skipped.
+                Signature: (provider, deployment, model_name, api_version, model_info) -> (action, reason)
+            include_model_id: If True, includes model_id in tuple (7 elements), else 6 elements.
+            include_load_balanced: If True, adds extra tests for load-balanced model groups.
+
+        Returns:
+            List of tuples with model test parameters.
+            - 6-element: (provider, deployment, model_name, api_version, action, reason)
+            - 7-element: (provider, deployment, model_name, api_version, model_id, action, reason)
+        """
+        models = cls._fetch_all_approved_models()
+        test_params: List[Tuple[str, ...]] = []
+        models_by_model_name: Dict[str, List[Tuple[str, ...]]] = defaultdict(list)
+
+        for model_data in models:
+            model_info = model_data.get("model_info", {}) or {}
+
+            provider, deployment, model_name, api_version = cls.extract_model_params(
+                model_data,
+            )
+
+            model_test_action, model_test_action_reason = should_skip_model(
+                provider,
+                deployment,
+                model_name,
+                api_version,
+                model_info,
+            )
+
+            if model_test_action.applicable():
+                if include_model_id:
+                    model_id = str(model_info.get("id"))
+                    params_tuple: Tuple[str, ...] = (
+                        provider,
+                        deployment,
+                        model_name,
+                        api_version,
+                        model_id,
+                        model_test_action,
+                        model_test_action_reason,
+                    )
+                else:
+                    params_tuple = (
+                        provider,
+                        deployment,
+                        model_name,
+                        api_version,
+                        model_test_action,
+                        model_test_action_reason,
+                    )
+
+                test_params.append(params_tuple)
+
+                if include_load_balanced:
+                    models_by_model_name[model_name].append(params_tuple)
+
+        if include_load_balanced and include_model_id:
+            for load_balanced_model_name, deployments in models_by_model_name.items():
+                if len(deployments) <= 1:
+                    continue
+
+                first_deployment = deployments[0]
+                test_params.append(
+                    (
+                        first_deployment[0],  # provider
+                        "load_balanced",
+                        load_balanced_model_name,
+                        "load_balanced",
+                        load_balanced_model_name,  # model_id = model_name for LB
+                        first_deployment[5],  # model_test_action
+                        first_deployment[6],  # model_test_action_reason
+                    ),
+                )
+
+        return test_params
+
+
+class UserKeyTestMixin:
+    """Mixin for tests that need to create users and API keys."""
+
+    allowed_routes: list[str] = []
+
+    _base_url: str = None
+    _master_api_key: str = None
+    admin_client: httpx.Client = None
+
+    @classmethod
+    def setup_admin_client(cls):
+        cls._base_url = get_litellm_base_url()
+        cls._master_api_key = get_litellm_api_key()
+        verify_setting = (
+            False
+            if use_local_litellm()
+            else BaseLiteLLMIntegrationTest._get_ssl_verify_setting()
+        )
+        cls.admin_client = httpx.Client(base_url=cls._base_url, verify=verify_setting)
+
+    @classmethod
+    def teardown_admin_client(cls):
+        if cls.admin_client:
+            cls.admin_client.close()
+
+    @staticmethod
+    def unique_suffix() -> str:
+        return f"{time.strftime('%Y%m%d%H%M%S')}{int(time.time() * 1000) % 1000:03d}"
+
+    @classmethod
+    def create_user_and_key(cls, user_suffix: str) -> tuple[str, str, str]:
+        user_email = f"test-user-{user_suffix}-{cls.unique_suffix()}@test.com"
+        user_response = cls.admin_client.post(
+            "/user/new",
+            json={
+                "user_email": user_email,
+                "user_alias": user_email,
+                "user_role": "internal_user",
+                "auto_create_key": "false",
+            },
+            headers={
+                "Authorization": f"Bearer {cls._master_api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=30,
+        )
+        assert user_response.status_code == 200, (
+            f"Failed to create user: {user_response.status_code} - {user_response.text}"
+        )
+        user_id = user_response.json().get("user_id")
+
+        key_alias = user_email.replace("@", "-at-").replace(".", "-")
+        key_response = cls.admin_client.post(
+            "/key/generate",
+            json={
+                "user_id": user_id,
+                "key_alias": key_alias,
+                "allowed_routes": cls.allowed_routes,
+            },
+            headers={
+                "Authorization": f"Bearer {cls._master_api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=30,
+        )
+        assert key_response.status_code == 200, (
+            f"Failed to create key: {key_response.status_code} - {key_response.text}"
+        )
+        api_key = key_response.json().get("key")
+
+        print(f"Created user {user_email}")
+        return user_id, api_key, user_email
+
+    @classmethod
+    def create_user_key_and_client(
+        cls,
+        user_suffix: str,
+    ) -> tuple[str, str, str, openai.OpenAI]:
+        user_id, api_key, user_email = cls.create_user_and_key(user_suffix)
+        verify_setting = (
+            False
+            if use_local_litellm()
+            else BaseLiteLLMIntegrationTest._get_ssl_verify_setting()
+        )
+        client = openai.OpenAI(
+            base_url=cls._base_url,
+            api_key=api_key,
+            http_client=httpx.Client(verify=verify_setting),
+        )
+        return user_id, api_key, user_email, client
+
+    @classmethod
+    def create_key_and_client(
+        cls,
+        user_id: str,
+        key_suffix: str,
+    ) -> tuple[str, openai.OpenAI]:
+        key_alias = f"additional-key-{key_suffix}-{cls.unique_suffix()}"
+        key_response = cls.admin_client.post(
+            "/key/generate",
+            json={
+                "user_id": user_id,
+                "key_alias": key_alias,
+                "allowed_routes": cls.allowed_routes,
+            },
+            headers={
+                "Authorization": f"Bearer {cls._master_api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=30,
+        )
+        assert key_response.status_code == 200, (
+            f"Failed to create additional key: {key_response.status_code} - {key_response.text}"
+        )
+        api_key = key_response.json().get("key")
+        verify_setting = (
+            False
+            if use_local_litellm()
+            else BaseLiteLLMIntegrationTest._get_ssl_verify_setting()
+        )
+        client = openai.OpenAI(
+            base_url=cls._base_url,
+            api_key=api_key,
+            http_client=httpx.Client(verify=verify_setting),
+        )
+        print(f"Created additional key for user {user_id}")
+        return api_key, client

tests/proxy_e2e_azure_batches_tests/conftest.py

@@ -0,0 +1,311 @@
+"""
+Pytest configuration for Azure Batch E2E Tests.
+
+This conftest manages:
+1. Mock Azure Batch server (FastAPI on port 8090)
+2. LiteLLM proxy server (port 4000)
+3. PostgreSQL database setup
+"""
+
+import asyncio
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Generator
+
+import httpx
+import pytest
+
+_test_dir = Path(__file__).parent
+sys.path.insert(0, str(_test_dir.parent.parent))  # litellm root
+sys.path.insert(0, str(_test_dir))  # test directory for local imports
+
+LOG_DIR = _test_dir
+
+
+def pytest_configure(config):
+    """Ensure test directory is in Python path before collection."""
+    test_dir = Path(__file__).parent
+    if str(test_dir) not in sys.path:
+        sys.path.insert(0, str(test_dir))
+
+
+MOCK_SERVER_PORT = 8090
+MOCK_SERVER_URL = f"http://localhost:{MOCK_SERVER_PORT}"
+LITELLM_PROXY_PORT = 4000
+LITELLM_PROXY_URL = f"http://localhost:{LITELLM_PROXY_PORT}"
+DATABASE_URL = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm"
+
+
+def kill_process_on_port(port: int) -> None:
+    """Kill any process using the specified port."""
+    try:
+        result = subprocess.run(
+            ["lsof", "-ti", f":{port}"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.stdout.strip():
+            pids = result.stdout.strip().split("\n")
+            for pid in pids:
+                try:
+                    subprocess.run(["kill", "-9", pid.strip()], timeout=5)
+                except Exception:
+                    pass
+            time.sleep(1)
+    except Exception:
+        pass
+
+
+def wait_for_server(url: str, max_attempts: int = 30, delay: float = 1.0) -> bool:
+    """Wait for a server to become available at url/health.
+
+    Any HTTP response (including 401) means the server is up.
+    Only connection errors count as "not ready yet".
+    """
+    for attempt in range(max_attempts):
+        try:
+            response = httpx.get(f"{url}/health", timeout=2.0)
+            return True
+        except (httpx.ConnectError, httpx.TimeoutException, httpx.NetworkError):
+            pass
+        except Exception:
+            pass
+        if attempt < max_attempts - 1:
+            time.sleep(delay)
+    return False
+
+
+def _read_log_tail(log_path: Path, max_lines: int = 80) -> str:
+    """Read the last N lines of a log file, returning empty string if not found."""
+    if not log_path.exists():
+        return "(log file not found)"
+    try:
+        text = log_path.read_text()
+        lines = text.strip().splitlines()
+        if len(lines) > max_lines:
+            return f"... ({len(lines) - max_lines} lines truncated) ...\n" + "\n".join(
+                lines[-max_lines:]
+            )
+        return text
+    except Exception as e:
+        return f"(error reading log: {e})"
+
+
+def _check_process_alive(process: subprocess.Popen, label: str, log_path: Path):
+    """Check if a subprocess crashed immediately after starting.
+    Raises pytest.fail with log output if the process has already exited.
+    """
+    time.sleep(1)
+    exit_code = process.poll()
+    if exit_code is not None:
+        log_output = _read_log_tail(log_path)
+        pytest.fail(
+            f"{label} exited immediately with code {exit_code}.\n"
+            f"--- {label} log ({log_path}) ---\n{log_output}\n"
+            f"--- end log ---"
+        )
+
+
+def setup_database() -> bool:
+    """Ensure PostgreSQL database exists and is accessible."""
+    try:
+        import psycopg2
+
+        conn = psycopg2.connect(
+            host="localhost",
+            port=5432,
+            database="litellm",
+            user="llmproxy",
+            password="dbpassword9090",
+            connect_timeout=5,
+        )
+        conn.close()
+        return True
+    except ImportError:
+        print("WARNING: psycopg2 not installed — cannot verify database")
+        return False
+    except Exception:
+        return False
+
+
+@pytest.fixture(scope="session")
+def mock_azure_server() -> Generator[str, None, None]:
+    """Start mock Azure batch server as a subprocess."""
+    print(f"\n{'=' * 60}")
+    print("Setting up Mock Azure Batch Server")
+    print(f"{'=' * 60}")
+
+    kill_process_on_port(MOCK_SERVER_PORT)
+
+    runner_script = Path(__file__).parent / "fixtures" / "run_mock_server.py"
+    runner_script.write_text(
+        """
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from fixtures.mock_azure_batch_server import create_mock_azure_batch_server
+import uvicorn
+
+if __name__ == "__main__":
+    app = create_mock_azure_batch_server()
+    uvicorn.run(app, host="0.0.0.0", port=8090, log_level="info", access_log=False)
+"""
+    )
+
+    mock_log = LOG_DIR / "mock_server.log"
+    log_file = open(mock_log, "w")
+
+    print(f"Starting mock server on port {MOCK_SERVER_PORT}...")
+    print(f"Log file: {mock_log}")
+    process = subprocess.Popen(
+        [sys.executable, str(runner_script)],
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        cwd=Path(__file__).parent,
+    )
+
+    _check_process_alive(process, "Mock server", mock_log)
+
+    if not wait_for_server(MOCK_SERVER_URL, max_attempts=30, delay=1.0):
+        log_output = _read_log_tail(mock_log)
+        exit_code = process.poll()
+        process.terminate()
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            process.wait()
+        log_file.close()
+        pytest.fail(
+            f"Mock server failed to start on port {MOCK_SERVER_PORT} "
+            f"(process exit_code={exit_code}).\n"
+            f"--- mock server log ---\n{log_output}\n--- end log ---\n"
+            f"Hint: ensure 'uvicorn' and 'fastapi' are installed."
+        )
+
+    print(f"Mock Azure server ready at {MOCK_SERVER_URL}")
+    yield MOCK_SERVER_URL
+
+    print("\nShutting down mock server...")
+    try:
+        process.terminate()
+        process.wait(timeout=5)
+    except subprocess.TimeoutExpired:
+        process.kill()
+        process.wait()
+    log_file.close()
+    print("Mock server stopped")
+
+
+@pytest.fixture(scope="session")
+def litellm_proxy_server(mock_azure_server: str) -> Generator[str, None, None]:
+    """Start LiteLLM proxy server for the test session."""
+    print(f"\n{'=' * 60}")
+    print("Setting up LiteLLM Proxy Server")
+    print(f"{'=' * 60}")
+
+    if not setup_database():
+        pytest.skip(
+            "PostgreSQL database not available at localhost:5432. "
+            "Start PostgreSQL and create a 'litellm' database:\n"
+            "  docker run -d --name litellm-db -p 5432:5432 "
+            '-e POSTGRES_USER=llmproxy -e POSTGRES_PASSWORD=dbpassword9090 '
+            "-e POSTGRES_DB=litellm postgres:15\n"
+            "Then run: prisma db push --schema=litellm/proxy/schema.prisma"
+        )
+    print("Database connection verified")
+
+    config_path = Path(__file__).parent / "fixtures" / "config.yml"
+    if not config_path.exists():
+        pytest.fail(f"Config file not found: {config_path}")
+    print("Config file found")
+
+    kill_process_on_port(LITELLM_PROXY_PORT)
+
+    os.environ["MOCK_SERVER_URL_V1"] = f"{mock_azure_server}/v1"
+    os.environ["MOCK_SERVER_URL_OPENAI_V1"] = f"{mock_azure_server}/openai/v1"
+    os.environ["DATABASE_URL"] = DATABASE_URL
+    os.environ["USE_LOCAL_LITELLM"] = "true"
+    os.environ["USE_MOCK_MODELS"] = "true"
+    os.environ["USE_STATE_TRACKER"] = "true"
+    os.environ["PROXY_BATCH_POLLING_INTERVAL"] = "10"
+
+    print("Environment configured")
+
+    print(f"Starting LiteLLM proxy on port {LITELLM_PROXY_PORT}...")
+    litellm_root = Path(__file__).parent.parent.parent
+
+    cmd = [
+        sys.executable,
+        "-m",
+        "litellm.proxy.proxy_cli",
+        "--config",
+        str(config_path),
+        "--port",
+        str(LITELLM_PROXY_PORT),
+        "--detailed_debug",
+    ]
+
+    proxy_log = LOG_DIR / "proxy_server.log"
+    log_file = open(proxy_log, "w")
+    print(f"Log file: {proxy_log}")
+
+    process = subprocess.Popen(
+        cmd,
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        env=os.environ.copy(),
+        cwd=litellm_root,
+    )
+
+    _check_process_alive(process, "LiteLLM proxy", proxy_log)
+
+    if not wait_for_server(LITELLM_PROXY_URL, max_attempts=60, delay=1.0):
+        log_output = _read_log_tail(proxy_log)
+        exit_code = process.poll()
+        process.terminate()
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            process.wait()
+        log_file.close()
+        pytest.fail(
+            f"LiteLLM proxy failed to start on port {LITELLM_PROXY_PORT} "
+            f"(process exit_code={exit_code}).\n"
+            f"--- proxy log (last 80 lines) ---\n{log_output}\n--- end log ---\n"
+            f"Hints:\n"
+            f"  1. Ensure Prisma client is generated: "
+            f"cd {litellm_root} && prisma generate --schema=litellm/proxy/schema.prisma\n"
+            f"  2. Ensure DB migrations are applied: "
+            f"prisma db push --schema=litellm/proxy/schema.prisma\n"
+            f"  3. Check the full log at: {proxy_log}"
+        )
+
+    print(f"LiteLLM proxy ready at {LITELLM_PROXY_URL}")
+    yield LITELLM_PROXY_URL
+
+    print("\nShutting down LiteLLM proxy...")
+    try:
+        process.terminate()
+        process.wait(timeout=10)
+    except subprocess.TimeoutExpired:
+        process.kill()
+        process.wait()
+    log_file.close()
+    print("LiteLLM proxy stopped")
+
+
+@pytest.fixture(scope="session")
+def event_loop():
+    """Provide an event loop for async tests."""
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    yield loop
+    loop.close()

tests/proxy_e2e_azure_batches_tests/fixtures/config.yml

@@ -0,0 +1,56 @@
+model_list:
+  - model_name: openai-fake-gpt-3.5-turbo
+    litellm_params:
+      model: openai/openai-fake-gpt-3.5-turbo
+      api_base: os.environ/MOCK_SERVER_URL_V1
+      api_key: fake-key
+  - model_name: openai-fake-gpt-4
+    litellm_params:
+      model: openai/openai-fake-gpt-4
+      api_base: os.environ/MOCK_SERVER_URL_V1
+      api_key: fake-key
+  - model_name: openai-fake-gpt-4o
+    litellm_params:
+      model: openai/openai-fake-gpt-4o
+      api_base: os.environ/MOCK_SERVER_URL_V1
+      api_key: fake-key
+  - model_name: fake-text-embedding-3-small
+    litellm_params:
+      model: openai/fake-text-embedding-3-small
+      api_base: os.environ/MOCK_SERVER_URL_V1
+      api_key: fake-key
+  - model_name: o3-mini-batch-2025-01-31
+    litellm_params:
+      model: openai/o3-mini-batch-2025-01-31
+      api_base: os.environ/MOCK_SERVER_URL_OPENAI_V1
+      api_key: fake-key
+    model_info:
+      mode: batch
+  - model_name: azure-fake-gpt-5-batch-2025-08-07
+    litellm_params:
+      api_base: http://0.0.0.0:8090
+      api_key: fake-key
+      api_version: 2025-03-01-preview
+      base_model: azure/gpt-5
+      model: azure/gpt-5-mini
+      custom_llm_provider: azure
+
+general_settings:
+  master_key: sk-1234
+  database_url: os.environ/DATABASE_URL
+  proxy_batch_polling_interval: 10
+
+litellm_settings:
+  drop_params: true
+  set_verbose: true
+  json_logs: true
+  # S3 callback for batch completion logging (points to mock server)
+  callbacks: ["s3_v2"]
+  s3_callback_params:
+    s3_bucket_name: litellm-test-bucket
+    s3_region_name: us-east-1
+    s3_endpoint_url: http://0.0.0.0:8090
+    s3_aws_access_key_id: fake-key
+    s3_aws_secret_access_key: fake-secret
+    s3_use_ssl: false
+    s3_verify: false

tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/__init__.py

@@ -0,0 +1,3 @@
+from .server import create_mock_azure_batch_server
+
+__all__ = ["create_mock_azure_batch_server"]

tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_azure_batch.py

@@ -0,0 +1,517 @@
+import asyncio
+import io
+import json
+import logging
+import time
+import uuid
+from typing import Dict, List, Optional
+
+from fastapi import FastAPI, HTTPException, Query, Request, UploadFile
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class FileObject(BaseModel):
+    id: str
+    object: str = "file"
+    bytes: int
+    created_at: int
+    filename: str
+    purpose: str
+    status: str = "processed"
+    status_details: Optional[str] = None
+    expires_at: Optional[int] = None
+
+
+class BatchObject(BaseModel):
+    id: str
+    object: str = "batch"
+    endpoint: str
+    errors: Optional[Dict] = None
+    input_file_id: str
+    completion_window: str
+    status: str
+    output_file_id: Optional[str] = None
+    error_file_id: Optional[str] = None
+    created_at: int
+    in_progress_at: Optional[int] = None
+    expires_at: Optional[int] = None
+    finalizing_at: Optional[int] = None
+    completed_at: Optional[int] = None
+    failed_at: Optional[int] = None
+    expired_at: Optional[int] = None
+    cancelling_at: Optional[int] = None
+    cancelled_at: Optional[int] = None
+    request_counts: Optional[Dict[str, int]] = None
+    metadata: Optional[Dict] = None
+
+
+class BatchListResponse(BaseModel):
+    object: str = "list"
+    data: List[Dict]
+    first_id: Optional[str] = None
+    last_id: Optional[str] = None
+    has_more: bool = False
+
+
+file_storage: Dict[str, Dict] = {}
+batch_storage: Dict[str, BatchObject] = {}
+batch_results: Dict[str, List[Dict]] = {}
+
+PROCESSING_DELAY_SECONDS = float(1)
+VALIDATING_DELAY_SECONDS = float(3)
+
+
+async def process_batch(batch_id: str):
+    logger.info(f"Starting batch processing for {batch_id}")
+    try:
+        batch = batch_storage[batch_id]
+
+        await asyncio.sleep(VALIDATING_DELAY_SECONDS)
+        batch.status = "in_progress"
+        batch.in_progress_at = int(time.time())
+        logger.info(f"Batch {batch_id} status: in_progress")
+
+        await process_batch_requests(batch_id)
+        await asyncio.sleep(PROCESSING_DELAY_SECONDS)
+
+        batch.status = "finalizing"
+        batch.finalizing_at = int(time.time())
+        logger.info(f"Batch {batch_id} status: finalizing")
+        await asyncio.sleep(PROCESSING_DELAY_SECONDS)
+
+        await create_output_file(batch_id)
+
+        batch.status = "completed"
+        batch.completed_at = int(time.time())
+        logger.info(f"Batch {batch_id} status: completed")
+
+    except Exception as e:
+        logger.error(f"Batch {batch_id} failed: {e}")
+        batch = batch_storage[batch_id]
+        batch.status = "failed"
+        batch.failed_at = int(time.time())
+        batch.errors = {
+            "object": "list",
+            "data": [{"code": "processing_error", "message": str(e)}],
+        }
+
+
+async def process_batch_requests(batch_id: str):
+    batch = batch_storage[batch_id]
+    input_file = file_storage[batch.input_file_id]
+
+    requests = []
+    for line in input_file["content"].split("\n"):
+        if line.strip():
+            try:
+                requests.append(json.loads(line))
+            except json.JSONDecodeError as e:
+                logger.warning(f"Invalid JSON line in batch {batch_id}: {e}")
+
+    logger.info(f"Batch {batch_id} has {len(requests)} requests")
+
+    results = []
+    failed_count = 0
+    for req in requests:
+        result = await process_single_request(req)
+        if result.get("error"):
+            failed_count += 1
+        results.append(result)
+
+    batch_results[batch_id] = results
+    batch.request_counts = {
+        "total": len(requests),
+        "completed": len(results) - failed_count,
+        "failed": failed_count,
+    }
+
+
+async def process_single_request(request_data: Dict) -> Dict:
+    custom_id = request_data.get("custom_id")
+    url = request_data.get("url", "/v1/chat/completions")
+    body = request_data.get("body", {})
+
+    if "/chat/completions" in url:
+        response_body = {
+            "id": f"chatcmpl-{uuid.uuid4().hex}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": body.get("model", "gpt-4o"),
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": "Mock batch response."},
+                    "finish_reason": "stop",
+                },
+            ],
+            "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+        }
+        status_code = 200
+    else:
+        response_body = {"error": {"message": f"Unsupported endpoint: {url}"}}
+        status_code = 400
+
+    return {
+        "id": f"batch_req_{uuid.uuid4().hex[:12]}",
+        "custom_id": custom_id,
+        "response": {
+            "status_code": status_code,
+            "request_id": f"req_{uuid.uuid4().hex[:12]}",
+            "body": response_body,
+        },
+        "error": None,
+    }
+
+
+async def create_output_file(batch_id: str):
+    results = batch_results.get(batch_id, [])
+    output_lines = [json.dumps(result) for result in results]
+    output_content = "\n".join(output_lines)
+
+    output_file_id = f"file-batch-output-{uuid.uuid4().hex[:12]}"
+    file_storage[output_file_id] = {
+        "content": output_content,
+        "filename": f"batch_output_{batch_id}.jsonl",
+        "purpose": "batch_output",
+        "bytes": len(output_content.encode()),
+        "created_at": int(time.time()),
+    }
+
+    batch = batch_storage[batch_id]
+    batch.output_file_id = output_file_id
+    logger.info(f"Created output file {output_file_id} for batch {batch_id}")
+
+
+def validate_batch_input(content: str) -> tuple[bool, str, List[Dict]]:
+    requests = []
+    custom_ids = set()
+
+    lines = content.strip().split("\n")
+    if not lines or all(not line.strip() for line in lines):
+        return False, "empty_batch", []
+
+    for line_num, line in enumerate(lines, 1):
+        if not line.strip():
+            continue
+        try:
+            req = json.loads(line)
+        except json.JSONDecodeError:
+            return False, "invalid_json_line", []
+
+        for field in ["custom_id", "method", "url", "body"]:
+            if field not in req:
+                return False, "invalid_request", []
+
+        if req["custom_id"] in custom_ids:
+            return False, "duplicate_custom_id", []
+        custom_ids.add(req["custom_id"])
+
+        requests.append(req)
+
+    if len(requests) > 100000:
+        return False, "too_many_tasks", []
+
+    return True, "", requests
+
+
+def setup_batch_routes(app: FastAPI):
+    # Files endpoints (OpenAI and Azure paths)
+    @app.post("/openai/v1/files")
+    @app.post("/openai/files")
+    @app.post("/v1/files")
+    @app.post("/files")
+    async def create_file(request: Request):
+        form = await request.form()
+        logger.info(f"File upload form fields: {list(form.keys())}")
+
+        file: UploadFile = form.get("file")
+        purpose: str = form.get("purpose", "batch")
+
+        if not file:
+            raise HTTPException(status_code=400, detail="No file provided")
+
+        logger.info(f"Uploading file: {file.filename}, purpose: {purpose}")
+
+        content = await file.read()
+        content_str = content.decode("utf-8")
+
+        file_id = f"file-{uuid.uuid4().hex[:24]}"
+        created_at = int(time.time())
+
+        expires_at = None
+        expires_after_seconds = form.get("expires_after[seconds]")
+        if expires_after_seconds:
+            try:
+                seconds = int(expires_after_seconds)
+                logger.info(f"expires_after[seconds] = {seconds}")
+                if seconds < 259200 or seconds > 2592000:
+                    raise HTTPException(
+                        status_code=400,
+                        detail={
+                            "error": {
+                                "code": "invalidPayload",
+                                "message": "Value for Seconds must be between 259200 and 2592000.",
+                            },
+                        },
+                    )
+                expires_at = created_at + seconds
+                logger.info(f"Calculated expires_at: {expires_at}")
+            except ValueError as e:
+                logger.warning(f"Failed to parse expires_after[seconds]: {e}")
+
+        file_storage[file_id] = {
+            "content": content_str,
+            "filename": file.filename or "batch_input.jsonl",
+            "purpose": purpose,
+            "bytes": len(content),
+            "created_at": created_at,
+            "expires_at": expires_at,
+        }
+
+        logger.info(f"Created file {file_id}, expires_at={expires_at}")
+        return FileObject(
+            id=file_id,
+            bytes=len(content),
+            created_at=created_at,
+            filename=file.filename or "batch_input.jsonl",
+            purpose=purpose,
+            expires_at=expires_at,
+        ).model_dump()
+
+    @app.get("/openai/v1/files/{file_id}")
+    @app.get("/openai/files/{file_id}")
+    @app.get("/v1/files/{file_id}")
+    @app.get("/files/{file_id}")
+    async def get_file(file_id: str):
+        logger.info(f"Getting file: {file_id}")
+        if file_id not in file_storage:
+            raise HTTPException(status_code=404, detail="File not found")
+
+        file_data = file_storage[file_id]
+        return FileObject(
+            id=file_id,
+            bytes=file_data["bytes"],
+            created_at=file_data["created_at"],
+            filename=file_data["filename"],
+            purpose=file_data["purpose"],
+            expires_at=file_data.get("expires_at"),
+        ).model_dump()
+
+    @app.get("/openai/v1/files/{file_id}/content")
+    @app.get("/openai/files/{file_id}/content")
+    @app.get("/v1/files/{file_id}/content")
+    @app.get("/files/{file_id}/content")
+    async def get_file_content(file_id: str):
+        logger.info(f"Getting file content: {file_id}")
+        if file_id not in file_storage:
+            raise HTTPException(status_code=404, detail="File not found")
+
+        file_data = file_storage[file_id]
+        content = file_data["content"]
+
+        return StreamingResponse(
+            io.StringIO(content),
+            media_type="application/octet-stream",
+            headers={
+                "Content-Disposition": f"attachment; filename={file_data['filename']}",
+            },
+        )
+
+    @app.delete("/openai/v1/files/{file_id}")
+    @app.delete("/openai/files/{file_id}")
+    @app.delete("/v1/files/{file_id}")
+    @app.delete("/files/{file_id}")
+    async def delete_file(file_id: str):
+        logger.info(f"Deleting file: {file_id}")
+        if file_id not in file_storage:
+            raise HTTPException(status_code=404, detail="File not found")
+
+        del file_storage[file_id]
+        return {"id": file_id, "object": "file", "deleted": True}
+
+    @app.get("/openai/v1/files")
+    @app.get("/openai/files")
+    @app.get("/v1/files")
+    @app.get("/files")
+    async def list_files(
+        purpose: Optional[str] = None,
+        limit: int = Query(10000, le=10000),
+    ):
+        logger.info(f"Listing files, purpose: {purpose}, limit: {limit}")
+        files = []
+        for file_id, file_data in file_storage.items():
+            if purpose is None or file_data.get("purpose") == purpose:
+                files.append(
+                    FileObject(
+                        id=file_id,
+                        bytes=file_data["bytes"],
+                        created_at=file_data["created_at"],
+                        filename=file_data["filename"],
+                        purpose=file_data["purpose"],
+                        expires_at=file_data.get("expires_at"),
+                    ).model_dump(),
+                )
+        return {"object": "list", "data": files[:limit]}
+
+    # Batches endpoints (OpenAI and Azure paths)
+    @app.post("/openai/v1/batches")
+    @app.post("/openai/batches")
+    @app.post("/v1/batches")
+    @app.post("/batches")
+    async def create_batch(request_data: dict):
+        input_file_id = request_data.get("input_file_id")
+        endpoint = request_data.get("endpoint", "/v1/chat/completions")
+        completion_window = request_data.get("completion_window", "24h")
+        metadata = request_data.get("metadata", {})
+        output_expires_after = request_data.get("output_expires_after")
+
+        logger.info(
+            f"Creating batch with input_file: {input_file_id}, endpoint: {endpoint}, output_expires_after: {output_expires_after}",
+        )
+
+        if not input_file_id or input_file_id not in file_storage:
+            raise HTTPException(status_code=400, detail="Input file not found")
+
+        input_file = file_storage[input_file_id]
+        is_valid, error_code, _ = validate_batch_input(input_file["content"])
+        if not is_valid:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": {
+                        "code": error_code,
+                        "message": f"Validation failed: {error_code}",
+                    },
+                },
+            )
+
+        batch_id = f"batch_{uuid.uuid4()}"
+        created_at = int(time.time())
+
+        if output_expires_after:
+            seconds = (
+                output_expires_after.get("seconds", 0)
+                if isinstance(output_expires_after, dict)
+                else 0
+            )
+            expires_at = created_at + seconds
+            logger.info(
+                f"Using output_expires_after: {seconds}s, expires_at: {expires_at}",
+            )
+        elif completion_window == "24h":
+            expires_at = created_at + (24 * 60 * 60)
+        else:
+            expires_at = created_at + (24 * 60 * 60)
+
+        batch = BatchObject(
+            id=batch_id,
+            endpoint=endpoint,
+            input_file_id=input_file_id,
+            completion_window=completion_window,
+            status="validating",
+            created_at=created_at,
+            expires_at=expires_at,
+            request_counts={"total": 0, "completed": 0, "failed": 0},
+            metadata=metadata,
+        )
+
+        batch_storage[batch_id] = batch
+        logger.info(f"Created batch {batch_id}")
+
+        asyncio.create_task(process_batch(batch_id))
+
+        return batch.model_dump()
+
+    @app.get("/openai/v1/batches/{batch_id}")
+    @app.get("/openai/batches/{batch_id}")
+    @app.get("/v1/batches/{batch_id}")
+    @app.get("/batches/{batch_id}")
+    async def get_batch(batch_id: str):
+        logger.info(f"Getting batch: {batch_id}")
+        if batch_id not in batch_storage:
+            raise HTTPException(status_code=404, detail="Batch not found")
+
+        return batch_storage[batch_id].model_dump()
+
+    @app.get("/openai/v1/batches")
+    @app.get("/openai/batches")
+    @app.get("/v1/batches")
+    @app.get("/batches")
+    async def list_batches(
+        after: Optional[str] = Query(None),
+        limit: int = Query(20, le=100),
+    ):
+        logger.info(f"Listing batches, after: {after}, limit: {limit}")
+        batches = list(batch_storage.values())
+        batches.sort(key=lambda x: x.created_at, reverse=True)
+
+        if after:
+            after_index = next((i for i, b in enumerate(batches) if b.id == after), -1)
+            if after_index >= 0:
+                batches = batches[after_index + 1 :]
+
+        batches = batches[:limit]
+
+        return BatchListResponse(
+            data=[batch.model_dump() for batch in batches],
+            first_id=batches[0].id if batches else None,
+            last_id=batches[-1].id if batches else None,
+            has_more=len(batches) == limit,
+        ).model_dump()
+
+    @app.post("/openai/v1/batches/{batch_id}/cancel")
+    @app.post("/openai/batches/{batch_id}/cancel")
+    @app.post("/v1/batches/{batch_id}/cancel")
+    @app.post("/batches/{batch_id}/cancel")
+    async def cancel_batch(batch_id: str):
+        logger.info(f"Cancelling batch: {batch_id}")
+        if batch_id not in batch_storage:
+            raise HTTPException(status_code=404, detail="Batch not found")
+
+        batch = batch_storage[batch_id]
+        if batch.status in ["completed", "failed", "cancelled", "expired"]:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Cannot cancel batch in {batch.status} status",
+            )
+
+        batch.status = "cancelled"
+        batch.cancelled_at = int(time.time())
+        logger.info(f"Batch {batch_id} cancelled")
+
+        return batch.model_dump()
+
+    # Debug endpoints
+    @app.get("/debug/batches")
+    async def debug_list_batches():
+        return {
+            "batches": {
+                batch_id: batch.model_dump()
+                for batch_id, batch in batch_storage.items()
+            },
+            "files": {
+                file_id: {k: v for k, v in data.items() if k != "content"}
+                for file_id, data in file_storage.items()
+            },
+        }
+
+    @app.post("/reset")
+    @app.post("/debug/clear")
+    async def reset_all():
+        file_storage.clear()
+        batch_storage.clear()
+        batch_results.clear()
+        logger.info("All data cleared")
+        return {"message": "All data cleared"}
+
+    @app.get("/debug/status")
+    async def debug_status():
+        return {
+            "files_count": len(file_storage),
+            "batches_count": len(batch_storage),
+            "batch_statuses": {bid: b.status for bid, b in batch_storage.items()},
+        }

tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_chat.py

@@ -0,0 +1,124 @@
+import json
+import time
+import uuid
+from datetime import datetime
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+
+
+def get_request_details(request: Request, body: dict = None) -> str:
+    details = {
+        "method": request.method,
+        "url": str(request.url),
+        "path": request.url.path,
+        "headers": dict(request.headers),
+        "query_params": dict(request.query_params),
+    }
+    return json.dumps(details, indent=2)
+
+
+def data_generator(response_details: str, model: str):
+    response_id = uuid.uuid4().hex
+    content = response_details
+    chunk_size = 50
+    for i in range(0, len(content), chunk_size):
+        text_chunk = content[i : i + chunk_size]
+        chunk = {
+            "id": f"chatcmpl-{response_id}",
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": model,
+            "choices": [{"index": 0, "delta": {"content": text_chunk}}],
+        }
+        yield f"data: {json.dumps(chunk)}\n\n"
+    final_chunk = {
+        "id": f"chatcmpl-{response_id}",
+        "object": "chat.completion.chunk",
+        "created": int(time.time()),
+        "model": model,
+        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+    }
+    yield f"data: {json.dumps(final_chunk)}\n\n"
+    yield "data: [DONE]\n\n"
+
+
+def setup_chat_routes(app: FastAPI):
+    @app.post("/chat/completions")
+    @app.post("/v1/chat/completions")
+    @app.post("/openai/deployments/{model:path}/chat/completions")
+    async def completion(request: Request):
+        data = await request.json()
+        model = data.get("model", "unknown")
+        request_details = get_request_details(request, data)
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        response_details = f"Request:{request_details}, Canned Response:{timestamp}"
+
+        if data.get("stream"):
+            return StreamingResponse(
+                content=data_generator(response_details, model),
+                media_type="text/event-stream",
+            )
+        else:
+            response_id = uuid.uuid4().hex
+            response = {
+                "id": f"chatcmpl-{response_id}",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": model,
+                "system_fingerprint": "fp_mock_server",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": response_details,
+                        },
+                        "logprobs": None,
+                        "finish_reason": "stop",
+                    },
+                ],
+                "usage": {
+                    "prompt_tokens": 9,
+                    "completion_tokens": 12,
+                    "total_tokens": 21,
+                },
+            }
+            return response
+
+    @app.post("/completions")
+    @app.post("/v1/completions")
+    async def text_completion(request: Request):
+        data = await request.json()
+        model = data.get("model", "unknown")
+        request_details = get_request_details(request, data)
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        response_details = f"Request:{request_details}, Canned Response:{timestamp}"
+
+        if data.get("stream"):
+            return StreamingResponse(
+                content=data_generator(response_details, model),
+                media_type="text/event-stream",
+            )
+        else:
+            response = {
+                "id": f"cmpl-{uuid.uuid4().hex}",
+                "choices": [
+                    {
+                        "finish_reason": "stop",
+                        "index": 0,
+                        "logprobs": None,
+                        "text": response_details,
+                    },
+                ],
+                "created": int(time.time()),
+                "model": model,
+                "object": "text_completion",
+                "system_fingerprint": None,
+                "usage": {
+                    "completion_tokens": 16,
+                    "prompt_tokens": 10,
+                    "total_tokens": 26,
+                },
+            }
+            return response

tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_embeddings.py

@@ -0,0 +1,23 @@
+from fastapi import FastAPI, Request
+
+
+def setup_embeddings_routes(app: FastAPI):
+    @app.post("/embeddings")
+    @app.post("/v1/embeddings")
+    @app.post("/openai/deployments/{model:path}/embeddings")
+    async def embeddings(request: Request):
+        data = await request.json()
+        model = data.get("model", "unknown")
+        _small_embedding = [
+            -0.006929283495992422,
+            -0.005336422007530928,
+            -4.547132266452536e-05,
+            -0.024047505110502243,
+        ]
+        big_embedding = _small_embedding * 100
+        return {
+            "object": "list",
+            "data": [{"object": "embedding", "index": 0, "embedding": big_embedding}],
+            "model": model,
+            "usage": {"prompt_tokens": 5, "total_tokens": 5},
+        }

tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_responses.py

@@ -0,0 +1,170 @@
+import json
+import re
+import time
+import uuid
+from datetime import datetime
+
+from typing import Any
+
+from fastapi import FastAPI, Request, HTTPException
+
+
+# Header to identify which model/deployment this request targets (simulates Azure model-specific encryption).
+# When set, the mock validates that encrypted_content in input was produced by this model.
+MOCK_AZURE_MODEL_HEADER = "X-Mock-Azure-Model"
+
+# Prefix we use in mock encrypted_content: gAAA_model_<model_id>_<32hex uuid>
+# Model id can contain underscores (e.g. gpt-5.1-codex-openai-2).
+ENCRYPTED_CONTENT_MODEL_PREFIX = re.compile(r"^gAAA_model_(.+)_[0-9a-f]{32}$")
+
+
+def _extract_model_from_encrypted_content(encrypted: str) -> str | None:
+    """Extract model id from our mock encrypted_content format, or None if not our format."""
+    if not isinstance(encrypted, str) or not encrypted.startswith("gAAA"):
+        return None
+    m = ENCRYPTED_CONTENT_MODEL_PREFIX.match(encrypted)
+    return m.group(1) if m else None
+
+
+def _collect_encrypted_contents(obj, out: list[str]) -> None:
+    """Recursively collect all encrypted_content string values from input structure."""
+    if isinstance(obj, dict):
+        if "encrypted_content" in obj and obj["encrypted_content"]:
+            out.append(obj["encrypted_content"])
+        for v in obj.values():
+            _collect_encrypted_contents(v, out)
+    elif isinstance(obj, list):
+        for item in obj:
+            _collect_encrypted_contents(item, out)
+
+
+def _validate_encrypted_content_model(request_model: str | None, input_data: Any) -> str | None:
+    """
+    If request_model is set, check that all encrypted_content in input was produced by this model.
+    Returns error message if validation fails, else None.
+    Content with our format (gAAA_model_<id>_) must match request_model.
+    """
+    if not request_model:
+        return None
+    encrypted_values: list[str] = []
+    _collect_encrypted_contents(input_data, encrypted_values)
+    for enc in encrypted_values:
+        content_model = _extract_model_from_encrypted_content(enc)
+        if content_model is not None and content_model != request_model:
+            err = enc[:50] + "..." if len(enc) > 50 else enc
+            return f"The encrypted content {err} could not be verified."
+    return None
+
+
+def get_request_details(request: Request, body: dict = None) -> str:
+    details = {
+        "method": request.method,
+        "url": str(request.url),
+        "path": request.url.path,
+        "headers": dict(request.headers),
+        "query_params": dict(request.query_params),
+    }
+    return json.dumps(details, indent=2)
+
+
+def setup_responses_routes(app: FastAPI):
+    @app.post("/responses")
+    @app.post("/v1/responses")
+    @app.post("/openai/responses")
+    async def responses_api(request: Request):
+        data = await request.json()
+        model = data.get("model", "unknown")
+
+        # Simulate Azure: encrypted content from one model cannot be verified by another.
+        input_data = data.get("input")
+        err_msg = _validate_encrypted_content_model(model, input_data)
+        if err_msg is not None:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": {
+                        "message": err_msg,
+                        "type": "invalid_request_error",
+                        "param": None,
+                        "code": "invalid_encrypted_content",
+                    }
+                },
+            )
+
+        request_details = get_request_details(request, data)
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        response_details = f"Request:{request_details}, Canned Response:{timestamp}"
+        response_id = uuid.uuid4().hex
+        message_id = f"msg_{uuid.uuid4().hex[:34]}"
+        reasoning_id = f"rs_{uuid.uuid4().hex[:34]}"
+
+        output_items: list[dict[str, Any]] = [
+            {
+                "id": message_id,
+                "content": [
+                    {
+                        "annotations": [],
+                        "text": response_details,
+                        "type": "output_text",
+                        "logprobs": [],
+                    },
+                ],
+                "role": "assistant",
+                "status": "completed",
+                "type": "message",
+            },
+        ]
+
+        if model:
+            output_items.append(
+                {
+                    "id": reasoning_id,
+                    "type": "reasoning",
+                    "status": "completed",
+                    "encrypted_content": f"gAAA_model_{model}_{uuid.uuid4().hex}",
+                }
+            )
+
+        return {
+            "id": f"resp_{response_id}",
+            "created_at": int(time.time()),
+            "error": None,
+            "incomplete_details": None,
+            "instructions": None,
+            "metadata": {},
+            "model": model,
+            "object": "response",
+            "output": output_items,
+            "parallel_tool_calls": True,
+            "temperature": data.get("temperature", 1.0),
+            "tool_choice": data.get("tool_choice", "auto"),
+            "tools": data.get("tools", []),
+            "top_p": data.get("top_p", 1.0),
+            "max_output_tokens": data.get("max_output_tokens"),
+            "previous_response_id": None,
+            "reasoning": {"effort": None, "summary": None},
+            "status": "completed",
+            "text": {"format": {"type": "text"}, "verbosity": "medium"},
+            "truncation": "disabled",
+            "usage": {
+                "input_tokens": 11,
+                "input_tokens_details": {
+                    "audio_tokens": None,
+                    "cached_tokens": 0,
+                    "text_tokens": None,
+                },
+                "output_tokens": 19,
+                "output_tokens_details": {"reasoning_tokens": 0, "text_tokens": None},
+                "total_tokens": 30,
+                "cost": None,
+            },
+            "user": None,
+            "store": True,
+            "background": False,
+            "content_filters": None,
+            "max_tool_calls": None,
+            "prompt_cache_key": None,
+            "safety_identifier": None,
+            "service_tier": "default",
+            "top_logprobs": 0,
+        }

tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_s3_callback.py

@@ -0,0 +1,98 @@
+"""
+Mock S3 callback receiver for testing LiteLLM S3 callbacks.
+
+This module provides S3-compatible endpoints that capture callback data
+sent by LiteLLM's s3_v2 callback handler after batch completion.
+"""
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class S3CallbackRecord(BaseModel):
+    key: str
+    bucket: str
+    content: Dict[str, Any]
+    timestamp: int
+    content_type: Optional[str] = None
+
+
+callback_storage: List[S3CallbackRecord] = []
+
+
+def setup_s3_callback_routes(app: FastAPI):
+    @app.put("/{bucket}/{key:path}")
+    async def s3_put_object(bucket: str, key: str, request: Request):
+        content_type = request.headers.get("content-type", "application/json")
+        body = await request.body()
+
+        try:
+            content = json.loads(body.decode("utf-8"))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            content = {"raw": body.decode("utf-8", errors="replace")}
+
+        record = S3CallbackRecord(
+            key=key,
+            bucket=bucket,
+            content=content,
+            timestamp=int(time.time()),
+            content_type=content_type,
+        )
+        callback_storage.append(record)
+
+        logger.info(f"S3 callback received: bucket={bucket}, key={key}")
+        logger.debug(f"Callback content: {json.dumps(content, indent=2)[:500]}")
+
+        return {
+            "ETag": f'"{hash(body)}"',
+            "VersionId": None,
+        }
+
+    @app.get("/mock-s3/callbacks")
+    async def list_callbacks(
+        bucket: Optional[str] = None,
+        key_prefix: Optional[str] = None,
+        limit: int = 100,
+    ):
+        results = callback_storage
+
+        if bucket:
+            results = [r for r in results if r.bucket == bucket]
+
+        if key_prefix:
+            results = [r for r in results if r.key.startswith(key_prefix)]
+
+        return {
+            "count": len(results),
+            "callbacks": [r.model_dump() for r in results[-limit:]],
+        }
+
+    @app.get("/mock-s3/callbacks/count")
+    async def count_callbacks(bucket: Optional[str] = None):
+        if bucket:
+            count = sum(1 for r in callback_storage if r.bucket == bucket)
+        else:
+            count = len(callback_storage)
+
+        return {"count": count}
+
+    @app.get("/mock-s3/callbacks/latest")
+    async def get_latest_callback():
+        if not callback_storage:
+            return {"callback": None}
+        return {"callback": callback_storage[-1].model_dump()}
+
+    @app.delete("/mock-s3/callbacks")
+    async def clear_callbacks():
+        count = len(callback_storage)
+        callback_storage.clear()
+        logger.info(f"Cleared {count} S3 callbacks")
+        return {"cleared": count}

tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/server.py

@@ -0,0 +1,33 @@
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+
+from .mock_azure_batch import setup_batch_routes
+from .mock_chat import setup_chat_routes
+from .mock_embeddings import setup_embeddings_routes
+from .mock_responses import setup_responses_routes
+from .mock_s3_callback import setup_s3_callback_routes
+
+
+def create_mock_azure_batch_server() -> FastAPI:
+    """Create a FastAPI app that mocks Azure Batch API and S3 callbacks."""
+    app = FastAPI()
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    @app.get("/health")
+    async def health():
+        return {"status": "ok"}
+
+    setup_chat_routes(app)
+    setup_responses_routes(app)
+    setup_embeddings_routes(app)
+    setup_batch_routes(app)
+    setup_s3_callback_routes(app)
+
+    return app

tests/proxy_e2e_azure_batches_tests/fixtures/run_mock_server.py

@@ -0,0 +1,12 @@
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from fixtures.mock_azure_batch_server import create_mock_azure_batch_server
+import uvicorn
+
+if __name__ == "__main__":
+    app = create_mock_azure_batch_server()
+    uvicorn.run(app, host="0.0.0.0", port=8090, log_level="info", access_log=False)

tests/proxy_e2e_azure_batches_tests/test_fixtures_smoke.py

@@ -0,0 +1,41 @@
+"""
+Smoke test to verify fixtures start and stop correctly.
+Run this first to ensure the infrastructure works before running full E2E tests.
+"""
+
+import httpx
+import pytest
+
+
+pytestmark = pytest.mark.usefixtures("mock_azure_server", "litellm_proxy_server")
+
+
+def test_mock_server_health(mock_azure_server):
+    """Verify mock Azure server is running and healthy."""
+    response = httpx.get(f"{mock_azure_server}/health", timeout=5.0)
+    assert response.status_code == 200
+    assert response.json() == {"status": "ok"}
+    print(f"✓ Mock Azure server is healthy at {mock_azure_server}")
+
+
+def test_litellm_proxy_health(litellm_proxy_server):
+    """Verify LiteLLM proxy is running and healthy."""
+    response = httpx.get(f"{litellm_proxy_server}/health", timeout=5.0)
+    assert response.status_code == 200
+    print(f"✓ LiteLLM proxy is healthy at {litellm_proxy_server}")
+
+
+def test_litellm_proxy_model_list(litellm_proxy_server):
+    """Verify LiteLLM proxy can list models."""
+    response = httpx.get(
+        f"{litellm_proxy_server}/v1/models",
+        headers={"Authorization": "Bearer sk-1234"},
+        timeout=5.0,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert "data" in data
+    models = [m["id"] for m in data["data"]]
+    print(f"✓ LiteLLM proxy has {len(models)} models configured")
+    assert "azure-fake-gpt-5-batch-2025-08-07" in models
+    print(f"✓ Azure batch model is configured")

tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py

@@ -0,0 +1,323 @@
+import base64
+import os
+import sys
+import time
+import warnings
+
+import httpx
+import openai
+import pytest
+from tenacity import RetryError
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+from base_integration_test import (
+    get_mock_server_base_url,
+    model_id,
+    use_mock_models,
+    UserKeyTestMixin,
+)
+from test_managed_files_base import (
+    ManagedFilesBase,
+    MIN_EXPIRY_SECONDS,
+    get_batch_model_names,
+)
+
+MANAGED_FILE_ID_PREFIX = "litellm_proxy"
+
+pytestmark = [
+    pytest.mark.usefixtures("mock_azure_server", "litellm_proxy_server"),
+    pytest.mark.skipif(
+        os.environ.get("SKIP_E2E_TESTS", "false").lower() == "true",
+        reason="E2E tests disabled via SKIP_E2E_TESTS env var"
+    ),
+]
+
+
+def is_managed_id(file_id: str) -> bool:
+    """Check if a file ID is a base64-encoded LiteLLM managed/unified ID."""
+    try:
+        padded = file_id + "=" * (-len(file_id) % 4)
+        decoded = base64.urlsafe_b64decode(padded).decode()
+        return decoded.startswith(MANAGED_FILE_ID_PREFIX)
+    except Exception:
+        return False
+
+
+def assert_managed_id(file_id: str, label: str):
+    assert is_managed_id(file_id), f"{label} should be a managed ID, got raw: {file_id}"
+
+
+def wip_features_enabled() -> bool:
+    return os.environ.get("WIP_FEATURES", "").lower() == "true"
+
+
+class TestManagedFilesAPI(ManagedFilesBase, UserKeyTestMixin):
+    @classmethod
+    def setup_class(cls):
+        super().setup_class()
+        cls.setup_admin_client()
+
+    @classmethod
+    def teardown_class(cls):
+        cls.teardown_admin_client()
+
+    @pytest.fixture(autouse=True)
+    def setup_test(self):
+        print(
+            f"\nBase URL: {self.base_url}, Using mock models: {use_mock_models()}",
+        )
+        self.clear_s3_callbacks()
+
+        user_id, api_key, user_email, client = self.create_user_key_and_client(
+            "e2e-batch",
+        )
+        self.test_user_id = user_id
+        self.openai_client = client
+        print(f"Using user {user_email} (id={user_id})")
+
+    def _create_and_verify_batch_input_file(self, tmp_path, model_name):
+        request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
+
+        print("Creating batch input file...")
+        batch_input_file = self.create_batch_input_file(
+            self.openai_client,
+            request_file,
+            MIN_EXPIRY_SECONDS,
+            target_model_names=model_name,
+        )
+        print(f"Created batch input file: {self.shorten_id(batch_input_file.id)}")
+        assert_managed_id(batch_input_file.id, "batch_input_file.id")
+
+        print("Retrieving batch input file metadata...")
+        metadata = self.openai_client.files.retrieve(batch_input_file.id)
+        assert_managed_id(metadata.id, "files.retrieve(input).id")
+        assert metadata.id == batch_input_file.id, (
+            f"Input file ID mismatch: retrieve returned '{metadata.id}' but expected '{batch_input_file.id}'"
+        )
+        assert metadata.object == "file"
+        assert metadata.bytes > 0, "bytes not set"
+        assert metadata.filename == "modified_file.jsonl"
+        assert metadata.purpose == "batch"
+        assert metadata.status in ["uploaded", "processed", "error"]
+        assert metadata.created_at > 0
+        if wip_features_enabled():
+            assert metadata.expires_at > 0, "expires_at not set"
+        self.print_file_metadata(metadata, "Input file")
+
+        return batch_input_file
+
+    def _create_and_verify_batch(self, input_file_id):
+        print("\nCreating batch...")
+        batch = self.create_batch(
+            self.openai_client,
+            input_file_id,
+            MIN_EXPIRY_SECONDS,
+        )
+        print(f"Created batch: {self.shorten_id(batch.id)}")
+
+        assert batch.id, "No batch ID returned"
+        assert_managed_id(batch.id, "batch.id")
+        assert_managed_id(batch.input_file_id, "batch.input_file_id")
+        assert batch.input_file_id == input_file_id, "batch.input_file_id mismatch"
+        assert batch.status in ["validating", "in_progress", "finalizing", "completed"]
+        if not batch.expires_at:
+            warnings.warn("batch expires_at not set")
+        else:
+            assert batch.expires_at > 0
+        if not batch.endpoint:
+            warnings.warn("batch.endpoint empty - Azure API quirk, not a bug")
+        else:
+            assert batch.endpoint == "/v1/chat/completions"
+        assert batch.completion_window == "24h"
+        assert batch.created_at > 0
+        self.print_batch_metadata(batch)
+
+        return batch
+
+    def _list_batches(self, batch_id, model_name):
+        if not wip_features_enabled():
+            return
+        print("\nListing batches...")
+        try:
+            batches_list = self.wait_for_batch_list(
+                model_name,
+                max_seconds=30,
+                wait_seconds=5,
+            )
+            batch_ids = [b.id for b in (batches_list.data if batches_list else [])]
+            if batch_id not in batch_ids:
+                warnings.warn(
+                    f"Batch {batch_id} not found in list. "
+                    f"batches.list returns raw IDs, not encoded IDs. raw IDs: {batch_ids}",
+                )
+        except openai.APIError as e:
+            pytest.fail(f"batches.list() failed: {e}")
+
+    def _wait_for_batch_completion(self, batch_id, tracker):
+        print(f"\nWaiting for batch {self.shorten_id(batch_id)} to complete...")
+        try:
+            batch_response = self.wait_for_batch_state(
+                self.openai_client,
+                batch_id,
+                "completed",
+                max_seconds=25 * 60,
+                wait_seconds=15,
+                state_tracker=tracker,
+            )
+        except RetryError:
+            tracker.print_state("Timeout waiting for batch completion")
+            raise TimeoutError("Timed out waiting for batch to be in state: completed")
+
+        assert_managed_id(batch_response.id, "batch_response.id")
+        assert batch_response.id == batch_id, (
+            f"batch_response.id mismatch: got '{batch_response.id}' but expected '{batch_id}'"
+        )
+        assert_managed_id(batch_response.input_file_id, "batch_response.input_file_id")
+        assert_managed_id(
+            batch_response.output_file_id,
+            "batch_response.output_file_id",
+        )
+
+        return batch_response
+
+    def _get_and_verify_batch_output(self, output_file_id):
+        print("\nRetrieving batch output file metadata...")
+        metadata = self.openai_client.files.retrieve(output_file_id)
+        assert_managed_id(metadata.id, "files.retrieve(output_file_id).id")
+        assert metadata.id == output_file_id, (
+            f"Output file ID mismatch: retrieve returned '{metadata.id}' but expected '{output_file_id}'"
+        )
+        assert metadata.object == "file"
+        assert metadata.bytes > 0, "bytes not set"
+        assert metadata.filename, "filename not set"
+        assert metadata.purpose in ["batch_output", "batch"]
+        assert metadata.created_at > 0
+        self.print_file_metadata(metadata, "Output file")
+
+        print("\nFetching batch output file content...")
+        content = self.openai_client.files.content(output_file_id)
+        assert content.text, "No batch file content returned"
+        assert len(content.text) > 0, "Batch file content is empty"
+        print(f"Output file content ({len(content.text)} bytes):")
+        for line in content.text.strip().split("\n")[:3]:
+            print(f"\t{line}")
+
+        return metadata
+
+    def _delete_file(self, file_id, label, max_retries=6, retry_delay=10):
+        print(f"\nDeleting {label}: {self.shorten_id(file_id)}")
+        for attempt in range(max_retries):
+            try:
+                self.openai_client.files.delete(file_id)
+                return
+            except openai.BadRequestError as e:
+                if "batch_processed" in str(e) and attempt < max_retries - 1:
+                    print(
+                        f"  File still referenced by unprocessed batch, "
+                        f"retrying in {retry_delay}s ({attempt + 1}/{max_retries})"
+                    )
+                    time.sleep(retry_delay)
+                else:
+                    pytest.fail(f"files.delete({label}) failed: {e}")
+            except openai.APIError as e:
+                pytest.fail(f"files.delete({label}) failed: {e}")
+
+    def _verify_file_deleted(self, file_id, label):
+        print(f"Verifying {label} is deleted...")
+        try:
+            self.openai_client.files.content(file_id)
+            assert False, f"{label} {file_id} still accessible after deletion"
+        except openai.NotFoundError:
+            print(f"{label} correctly not accessible after deletion")
+
+    # ------------------------------------------------------------------
+    # Tests
+    # ------------------------------------------------------------------
+
+    @pytest.mark.parametrize(
+        "model_name",
+        get_batch_model_names(),
+        ids=model_id,
+    )
+    def test_e2e_managed_batch(self, tmp_path, model_name):
+        print(
+            f"\n\nStarting test with base_url={self.base_url} and model_name={model_name}\n",
+        )
+        self.reset_mock_server()
+        tracker = self.create_state_tracker()
+
+        batch_input_file = self._create_and_verify_batch_input_file(
+            tmp_path,
+            model_name,
+        )
+        tracker.set_file_id(batch_input_file.id)
+        tracker.print_state("After creating batch input file")
+
+        batch = self._create_and_verify_batch(batch_input_file.id)
+        tracker.set_batch_id(batch.id)
+        tracker.print_state("After creating batch")
+
+        self._list_batches(batch.id, model_name)
+
+        batch_response = self._wait_for_batch_completion(batch.id, tracker)
+        tracker.print_state("After batch completed")
+
+        self._get_and_verify_batch_output(batch_response.output_file_id)
+        tracker.print_state("After retrieving output file")
+
+        tracker.print_state("Final state after cleanup")
+        tracker.wait_and_print_s3_callbacks()
+        tracker.assert_batch_cost_callback()
+
+        self._delete_file(batch_input_file.id, "input file")
+        self._delete_file(batch_response.output_file_id, "output file")
+
+        self._verify_file_deleted(batch_input_file.id, "input file")
+        self._verify_file_deleted(batch_response.output_file_id, "output file")
+
+    def cleanup_batches_in_database(self):
+        import psycopg2
+
+        print("Cleaning up stale batch records from database...")
+        try:
+            conn = psycopg2.connect(
+                host="localhost",
+                port=5432,
+                database="litellm",
+                user="llmproxy",
+                password="dbpassword9090",
+            )
+            with conn.cursor() as cur:
+                cur.execute("""
+                    DELETE FROM "LiteLLM_ManagedObjectTable"
+                    WHERE file_purpose = 'batch' AND status = 'validating'
+                """)
+                deleted = cur.rowcount
+                conn.commit()
+                if deleted > 0:
+                    print(f"Deleted {deleted} stale batch records")
+            conn.close()
+        except Exception as e:
+            print(f"Warning: Could not clean up database: {e}")
+
+    def clear_s3_callbacks(self):
+        clear_response = httpx.delete(f"{get_mock_server_base_url()}/mock-s3/callbacks")
+        assert clear_response.status_code == 200, (
+            f"Failed to clear callbacks: {clear_response.text}"
+        )
+        return clear_response.json()
+
+    @pytest.mark.skipif(
+        True,
+        reason="Skipping managed files test till managed files feature is available",
+    )
+    @pytest.mark.parametrize(
+        "model_name",
+        get_batch_model_names(),
+        ids=model_id,
+    )
+    def test_error_files(self, tmp_path, model_name):
+        raise NotImplementedError(
+            "To implement. Fail a batch and retrieve the error file.",
+        )

tests/proxy_e2e_azure_batches_tests/validate_e2e_setup.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+"""
+Validation script for Azure Batch E2E test setup.
+Run this before running the actual tests to verify all components are accessible.
+"""
+
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+def check_imports():
+    """Verify all required imports work."""
+    print("Checking imports...")
+    try:
+        from base_integration_test import (
+            get_mock_server_base_url,
+            get_litellm_base_url,
+            get_litellm_api_key,
+        )
+        print("  ✓ base_integration_test imports OK")
+        
+        from test_managed_files_base import ManagedFilesBase, get_batch_model_names
+        print("  ✓ test_managed_files_base imports OK")
+        
+        from fixtures.mock_azure_batch_server import create_mock_azure_batch_server
+        print("  ✓ mock_azure_batch_server imports OK")
+        
+        import httpx
+        import openai
+        import psycopg2
+        import uvicorn
+        print("  ✓ All external dependencies OK")
+        
+        return True
+    except ImportError as e:
+        print(f"  ✗ Import error: {e}")
+        return False
+
+
+def check_config_file():
+    """Verify config file exists."""
+    print("\nChecking config file...")
+    config_path = Path(__file__).parent / "fixtures" / "config.yml"
+    if config_path.exists():
+        print(f"  ✓ Config file found: {config_path}")
+        return True
+    else:
+        print(f"  ✗ Config file not found: {config_path}")
+        return False
+
+
+def check_database():
+    """Verify database connection."""
+    print("\nChecking database connection...")
+    try:
+        import psycopg2
+        conn = psycopg2.connect(
+            host="localhost",
+            port=5432,
+            database="litellm",
+            user="llmproxy",
+            password="dbpassword9090",
+        )
+        conn.close()
+        print("  ✓ Database connection OK")
+        return True
+    except Exception as e:
+        print(f"  ✗ Database connection failed: {e}")
+        print("    Start PostgreSQL with:")
+        print("    docker run --name litellm-postgres -e POSTGRES_USER=llmproxy \\")
+        print("      -e POSTGRES_PASSWORD=dbpassword9090 -e POSTGRES_DB=litellm \\")
+        print("      -p 5432:5432 -d postgres:15")
+        return False
+
+
+def check_ports():
+    """Check if required ports are available."""
+    print("\nChecking ports...")
+    import socket
+    
+    for port, name in [(4000, "LiteLLM Proxy"), (8090, "Mock Server")]:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            try:
+                s.bind(("localhost", port))
+                print(f"  ✓ Port {port} ({name}) is available")
+            except OSError:
+                print(f"  ⚠ Port {port} ({name}) is in use (will reuse if healthy)")
+    return True
+
+
+def main():
+    print("=" * 70)
+    print("Azure Batch E2E Test Setup Validation")
+    print("=" * 70)
+    
+    checks = [
+        check_imports(),
+        check_config_file(),
+        check_database(),
+        check_ports(),
+    ]
+    
+    print("\n" + "=" * 70)
+    if all(checks):
+        print("✓ All checks passed! Ready to run E2E tests.")
+        print("\nRun tests with:")
+        print("  cd litellm")
+        print("  export DATABASE_URL='postgresql://llmproxy:dbpassword9090@localhost:5432/litellm'")
+        print("  poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py -vv")
+        return 0
+    else:
+        print("✗ Some checks failed. Please fix the issues above.")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())