litellm/tests/e2e/batches/test_batch_cleanup.py

from builtins import ExceptionGroup
from collections.abc import Callable
from typing import Final
from unittest.mock import Mock, call

import pytest
from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result
from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form
from capabilities import CAPABILITIES, Capability
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
from lifecycle import ResourceManager
from models import KeyGenerateBody

MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE="
MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x"


class ExpectedCalls[T]:
    def __init__(self, values: tuple[T, ...]) -> None:
        self.values: Final = values
        self.recorder: Final = Mock()

    def __call__(self, value: T) -> None:
        self.recorder(value)

    def assert_done(self) -> None:
        assert tuple(self.recorder.call_args_list) == tuple(call(value) for value in self.values)


class CleanupClient:
    def __init__(
        self,
        *,
        calls: ExpectedCalls[str],
        files: tuple[Result[FileDeleteResponse], ...] = (),
        batches: tuple[Result[BatchObject], ...] = (),
        cancellations: tuple[Result[BatchObject], ...] = (),
    ) -> None:
        self.calls: Final = calls
        self.file_response: Final[Callable[[], Result[FileDeleteResponse]]] = Mock(side_effect=files)
        self.batch_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=batches)
        self.cancel_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=cancellations)

    def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]:
        self.calls(f"delete {provider} {file_id}")
        return self.file_response()

    def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
        self.calls(f"retrieve {provider} {batch_id}")
        return self.batch_response()

    def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
        self.calls(f"cancel {provider} {batch_id}")
        return self.cancel_response()

    def generate_key(self, body: KeyGenerateBody) -> str:
        return "test-key"

    def delete_key(self, key: str) -> None:
        self.calls(f"delete key {key}")

    def delete_customers(self, user_ids: list[str]) -> None:
        self.calls(f"delete customers {user_ids}")


def batch(status: str) -> Success[BatchObject]:
    return Success(status_code=200, data=BatchObject(id="batch-1", status=status))


def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]:
    return Success(status_code=200, data=FileDeleteResponse(id="file-1", deleted=deleted))


class TestFileCleanup:
    def test_managed_delete_accepts_the_deleted_file_object(self) -> None:
        response: Final = Success(
            status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"})
        )
        client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(response,))
        cleanup_file(client, MANAGED_FILE_ID, key="test-key")
        client.calls.assert_done()

    @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID])
    def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls((f"delete None {file_id}",)),
            files=(Success(status_code=200, data=FileDeleteResponse(id=file_id)),),
        )
        with pytest.raises(AssertionError, match="did not confirm deletion"):
            cleanup_file(client, file_id, key="test-key")
        client.calls.assert_done()

    @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES])
    def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None:
        expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None
        client: Final = CleanupClient(
            calls=ExpectedCalls((f"delete {expected_provider} file-1",)), files=(deleted_file(),)
        )
        cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider)
        client.calls.assert_done()

    def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(("delete azure file-1", "delete key test-key")),
            files=(UnknownApiError(status_code=403, body="secret response"),),
        )
        manager: Final = ResourceManager(client=client, strict_cleanup=True)
        key: Final = manager.key()
        manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure"))
        with pytest.raises(ExceptionGroup) as caught:
            manager.teardown()
        client.calls.assert_done()
        assert len(caught.value.exceptions) == 1
        assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403"

    def test_success_response_must_confirm_deletion(self) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(("delete None file-1",)), files=(deleted_file(deleted=False),)
        )
        with pytest.raises(AssertionError, match="did not confirm deletion"):
            cleanup_file(client, "file-1", key="test-key")
        client.calls.assert_done()

    def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(("delete azure file-1",)),
            files=(UnknownApiError(status_code=404, body="missing"),),
        )
        cleanup_file(client, "file-1", key="test-key", provider="azure")
        client.calls.assert_done()

    def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(("delete None file-1", "delete key test-key")),
            files=(UnknownApiError(status_code=403, body="forbidden"),),
        )
        manager: Final = ResourceManager(client=client)
        key: Final = manager.key()
        manager.defer(lambda: cleanup_file(client, "file-1", key=key))
        manager.teardown()
        client.calls.assert_done()


class TestCleanupRetries:
    @pytest.mark.parametrize(
        "failure",
        [NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")],
    )
    def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None:
        responses: Final = (failure, deleted_file())
        outcomes: Final = Mock(side_effect=responses)
        delays: Final = ExpectedCalls((1.0,))
        result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays)
        assert isinstance(result, Success) and result.data.deleted
        delays.assert_done()

    def test_persistent_error_has_bounded_retries(self) -> None:
        failure: Final = UnknownApiError(status_code=503, body="unavailable")
        outcomes: Final = Mock(return_value=failure)
        delays: Final = ExpectedCalls(CLEANUP_DELAYS)
        result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays)
        assert result is failure
        delays.assert_done()
        assert outcomes.call_count == len(CLEANUP_DELAYS) + 1

    def test_permanent_error_is_not_retried(self) -> None:
        failure: Final = UnknownApiError(status_code=403, body="forbidden")
        responses: Final = (failure, deleted_file())
        outcomes: Final = Mock(side_effect=responses)
        delays: Final = ExpectedCalls[float](())
        assert cleanup_result(outcomes, wait=delays) is failure
        delays.assert_done()
        assert outcomes.call_count == 1


class TestBatchCancellation:
    def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3),
            batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")),
        )
        delays: Final = ExpectedCalls((10.0,))
        cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays)
        client.calls.assert_done()
        delays.assert_done()

    def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(
                (
                    f"retrieve None {MANAGED_BATCH_ID}",
                    f"retrieve None {MANAGED_BATCH_ID}",
                    "delete None file-1",
                    "delete key test-key",
                )
            ),
            batches=(batch("cancelling"), batch("cancelling")),
            files=(deleted_file(),),
        )
        times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS)
        ticks: Final[Callable[[], float]] = Mock(side_effect=times)
        manager: Final = ResourceManager(client=client, strict_cleanup=True)
        key: Final = manager.key()
        manager.defer(lambda: cleanup_file(client, "file-1", key=key))
        manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks))
        with pytest.raises(ExceptionGroup) as caught:
            manager.teardown()
        assert "cancellation did not finish" in str(caught.value.exceptions[0])
        client.calls.assert_done()

    @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"])
    def test_inactive_batch_needs_no_cancellation(self, status: str) -> None:
        client: Final = CleanupClient(calls=ExpectedCalls(("retrieve None batch-1",)), batches=(batch(status),))
        cleanup_batch(client, "batch-1", key="test-key")
        client.calls.assert_done()

    def test_active_batch_is_cancelled_through_its_provider(self) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(("retrieve azure batch-1", "cancel azure batch-1")),
            batches=(batch("in_progress"), batch("cancelled")),
            cancellations=(batch("cancelling"),),
        )
        cleanup_batch(client, "batch-1", key="test-key", provider="azure")
        client.calls.assert_done()

    @pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID])
    @pytest.mark.parametrize("pending_status", ["validating", "in_progress"])
    def test_accepted_cancellation_waits_through_stale_provider_status(
        self, batch_id: str, pending_status: str
    ) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(
                (
                    f"retrieve vertex_ai {batch_id}",
                    f"cancel vertex_ai {batch_id}",
                    f"retrieve vertex_ai {batch_id}",
                    f"retrieve vertex_ai {batch_id}",
                    f"retrieve vertex_ai {batch_id}",
                    "delete vertex_ai file-1",
                    "delete key test-key",
                )
            ),
            batches=(batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled")),
            cancellations=(batch(pending_status),),
            files=(deleted_file(),),
        )
        delays: Final = ExpectedCalls((10.0, 10.0))
        manager: Final = ResourceManager(client=client, strict_cleanup=True)
        key: Final = manager.key()
        manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai"))
        manager.defer(lambda: cleanup_batch(client, batch_id, key=key, provider="vertex_ai", wait=delays))
        manager.teardown()
        client.calls.assert_done()
        delays.assert_done()

    @pytest.mark.parametrize("output_delete_fails", [False, True])
    def test_batch_that_completed_before_cleanup_deletes_output_and_error_files(
        self, output_delete_fails: bool
    ) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")),
            batches=(
                Success(
                    status_code=200,
                    data=BatchObject(
                        id="batch-1",
                        status="completed",
                        input_file_id="file-input",
                        output_file_id="file-output",
                        error_file_id="file-error",
                    ),
                ),
            ),
            files=(
                UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(),
                deleted_file(),
            ),
        )
        if output_delete_fails:
            with pytest.raises(ExceptionGroup, match="output cleanup failed"):
                cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
        else:
            cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
        client.calls.assert_done()

    @pytest.mark.parametrize("status", ["completed", "in_progress"])
    def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None:
        client: Final = CleanupClient(
            calls=ExpectedCalls(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1")),
            batches=(batch("in_progress"), batch(status)),
            cancellations=(UnknownApiError(status_code=409, body="conflict"),),
        )
        if status == "completed":
            cleanup_batch(client, "batch-1", key="test-key")
        else:
            with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"):
                cleanup_batch(client, "batch-1", key="test-key")
        client.calls.assert_done()


class TestAzureFileExpiry:
    def test_azure_form_serializes_native_expiry_for_the_proxy(self) -> None:
        form: Final = batch_upload_form("azure", target_model_names="azure-test")
        assert form.model_dump(by_alias=True, exclude_none=True) == {
            "purpose": "batch",
            "target_model_names": "azure-test",
            "expires_after[anchor]": "created_at",
            "expires_after[seconds]": AZURE_FILE_EXPIRY_SECONDS,
        }

    @pytest.mark.parametrize("provider", ["openai", "vertex_ai", "bedrock"])
    def test_other_providers_keep_their_existing_upload_fields(self, provider: str) -> None:
        assert batch_upload_form(provider).model_dump(by_alias=True, exclude_none=True) == {"purpose": "batch"}