litellm/tests/unit/llms/bedrock/batches/test_handler.py

"""Unit tests for ``BedrockBatchesHandler._handle_model_invocation_job_status``.

These cover the upstream support for retrieving Bedrock bulk batch jobs
(``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<id>``) — the ARN
type returned by ``CreateModelInvocationJob``. We mock the boto3 client so
the tests don't hit AWS.
"""

from __future__ import annotations

import json
from collections.abc import Iterator, Mapping
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import MagicMock, patch

import pytest
from botocore.awsrequest import AWSPreparedRequest, AWSResponse


from litellm.llms.bedrock.batches.handler import (  # noqa: E402
    BedrockBatchesHandler,
    _extract_job_id_from_arn,
    _extract_region_from_bedrock_arn,
    _predict_output_file_uri,
    _to_epoch,
)

JOB_ID = "abc1234567"
JOB_ARN = f"arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/{JOB_ID}"
INPUT_URI = "s3://my-bucket/inputs/qwen3-235b-a22b-2507-batch.jsonl"
OUTPUT_PREFIX = "s3://my-bucket/litellm-batch-outputs/litellm-bedrock-files-qwen-uuid/"
SUBMIT_TIME = datetime(2026, 4, 28, 12, 0, 0, tzinfo=timezone.utc)
END_TIME = datetime(2026, 4, 28, 12, 30, 0, tzinfo=timezone.utc)


def _fake_boto3_response(status: str = "Completed", end_time=END_TIME):
    return {
        "jobArn": JOB_ARN,
        "jobName": "litellm-bedrock-files-qwen-uuid",
        "modelId": "bedrock/qwen.qwen3-235b-a22b-2507-v1:0",
        "status": status,
        "submitTime": SUBMIT_TIME,
        "lastModifiedTime": end_time,
        "endTime": end_time,
        "inputDataConfig": {"s3InputDataConfig": {"s3Uri": INPUT_URI}},
        "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_PREFIX}},
    }


@pytest.fixture
def patched_boto3():
    """Yield a stub bedrock client whose `get_model_invocation_job` is a MagicMock."""
    fake_client = MagicMock()
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response()
    with (
        patch("boto3.client", return_value=fake_client) as boto_client_factory,
        patch(
            "litellm.llms.bedrock.batches.transformation.BedrockBatchesConfig.get_credentials",
            return_value=MagicMock(access_key="AKIA", secret_key="SECRET", token=None),
        ),
    ):
        yield fake_client, boto_client_factory


def test_extract_region_from_arn():
    assert _extract_region_from_bedrock_arn(JOB_ARN) == "us-west-2"
    assert _extract_region_from_bedrock_arn("arn:aws:bedrock::123:foo/bar") is None
    assert _extract_region_from_bedrock_arn("not-an-arn") is None


def test_extract_region_swallows_unexpected_split_errors():
    """Defensive `except Exception` branch — anything that isn't a plain str
    should fall through to ``None`` rather than blow up."""

    class WeirdArn:
        def split(self, _sep):
            raise RuntimeError("boom")

    assert _extract_region_from_bedrock_arn(WeirdArn()) is None  # type: ignore[arg-type]


def test_predict_output_file_uri_returns_none_for_directory_input_uri():
    """Input URI ending in `/` has an empty basename — we must bail rather
    than emit ``<prefix>/<job-id>/.out``."""
    assert (
        _predict_output_file_uri(OUTPUT_PREFIX, "s3://bucket/inputs/", JOB_ID) is None
    )


_DT = datetime(2026, 4, 28, 12, 0, 0, tzinfo=timezone.utc)


@pytest.mark.parametrize(
    "value,expected",
    [
        (None, None),
        (1730000000, 1730000000),
        (1730000000.5, 1730000000),
        (_DT, int(_DT.timestamp())),
        ("2026-04-28T12:00:00Z", None),  # strings aren't supported -> None
    ],
)
def test_to_epoch_handles_supported_types(value, expected):
    assert _to_epoch(value) == expected


def test_extract_job_id_from_arn():
    assert _extract_job_id_from_arn(JOB_ARN) == JOB_ID
    assert (
        _extract_job_id_from_arn("arn:aws:bedrock:us-west-2:1:async-invoke/x") is None
    )


def test_predict_output_file_uri_happy_path():
    expected = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
    assert _predict_output_file_uri(OUTPUT_PREFIX, INPUT_URI, JOB_ID) == expected


def test_predict_output_file_uri_adds_trailing_slash():
    prefix_no_slash = OUTPUT_PREFIX.rstrip("/")
    expected = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
    assert _predict_output_file_uri(prefix_no_slash, INPUT_URI, JOB_ID) == expected


@pytest.mark.parametrize(
    "missing_arg",
    [
        ("", INPUT_URI, JOB_ID),
        (OUTPUT_PREFIX, "", JOB_ID),
        (OUTPUT_PREFIX, INPUT_URI, None),
    ],
)
def test_predict_output_file_uri_returns_none_when_missing_input(missing_arg):
    assert _predict_output_file_uri(*missing_arg) is None


def test_handle_model_invocation_job_status_completed(patched_boto3):
    fake_client, boto_client_factory = patched_boto3

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    fake_client.get_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)

    # Region should be sniffed from the ARN.
    _, kwargs = boto_client_factory.call_args
    assert kwargs["region_name"] == "us-west-2"

    assert batch.id == JOB_ARN
    assert batch.status == "completed"
    assert batch.input_file_id == INPUT_URI
    expected_out = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
    assert batch.output_file_id == expected_out
    assert batch.completed_at == int(END_TIME.timestamp())
    assert batch.failed_at is None
    assert batch.cancelled_at is None
    assert batch.request_counts is None
    assert batch.metadata["job_arn"] == JOB_ARN
    assert batch.metadata["output_file_uri"] == expected_out
    assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX


@pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)])
def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count):
    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = {
        **_fake_boto3_response(),
        "totalRecordCount": 100,
        "successRecordCount": success_count,
        "errorRecordCount": error_count,
    }

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.request_counts is not None
    assert (batch.request_counts.total, batch.request_counts.completed, batch.request_counts.failed) == (
        100,
        success_count,
        error_count,
    )


def test_missing_record_counts_leave_request_counts_none(patched_boto3):
    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response()

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.request_counts is None


def test_total_without_success_count_leaves_request_counts_none(patched_boto3):
    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = {**_fake_boto3_response(), "totalRecordCount": 100}

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.request_counts is None


def test_missing_error_count_maps_to_zero_failed(patched_boto3):
    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = {
        **_fake_boto3_response(),
        "totalRecordCount": 100,
        "successRecordCount": 100,
    }

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.request_counts is not None
    assert (batch.request_counts.total, batch.request_counts.completed, batch.request_counts.failed) == (100, 100, 0)


@pytest.mark.parametrize(
    "bedrock_status,openai_status",
    [
        ("Submitted", "validating"),
        ("Validating", "validating"),
        ("Scheduled", "validating"),
        ("InProgress", "in_progress"),
        ("Stopping", "cancelling"),
        ("Stopped", "cancelled"),
        ("Completed", "completed"),
        ("PartiallyCompleted", "completed"),
        ("Failed", "failed"),
        ("Expired", "expired"),
        # Unknown/unmapped Bedrock status falls back to "in_progress" so we
        # don't 500 on a future AWS-side enum addition.
        ("MyBrandNewStatus", "in_progress"),
    ],
)
def test_status_mapping(patched_boto3, bedrock_status, openai_status):
    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(
        status=bedrock_status
    )

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.status == openai_status
    # output_file_id is only populated for terminal-completed jobs, so callers
    # don't accidentally try to download a non-existent file mid-run.
    if openai_status == "completed":
        assert batch.output_file_id is not None
    else:
        assert batch.output_file_id is None


def test_explicit_region_overrides_arn(patched_boto3):
    _, boto_client_factory = patched_boto3
    BedrockBatchesHandler._handle_model_invocation_job_status(
        batch_id=JOB_ARN, aws_region_name="eu-central-1"
    )
    _, kwargs = boto_client_factory.call_args
    assert kwargs["region_name"] == "eu-central-1"


def test_failure_message_propagates(patched_boto3):
    fake_client, _ = patched_boto3
    failed_response = _fake_boto3_response(status="Failed")
    failed_response["message"] = "Input file failed validation"
    fake_client.get_model_invocation_job.return_value = failed_response

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.status == "failed"
    assert batch.failed_at == int(END_TIME.timestamp())
    assert batch.metadata["failure_message"] == "Input file failed validation"


def test_completed_with_unpredictable_output_uri_stays_none(patched_boto3):
    """
    Regression guard for the original NoSuchKey bug: if Bedrock's response is
    missing pieces we need to compute the per-job output file path (here, the
    input s3Uri), `output_file_id` must stay `None` rather than fall back to
    the bare prefix. Falling back to the prefix is what produced the original
    NoSuchKey error this PR fixes.
    """
    fake_client, _ = patched_boto3
    incomplete_response = _fake_boto3_response(status="Completed")
    incomplete_response["inputDataConfig"] = {"s3InputDataConfig": {"s3Uri": ""}}
    fake_client.get_model_invocation_job.return_value = incomplete_response

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.status == "completed"
    # output_file_id MUST be None (not the bare prefix) — that's the whole
    # point of this regression test. Callers branch on this field.
    assert batch.output_file_id is None
    # The metadata field uses "" because OpenAI Batch metadata is dict[str, str];
    # callers should branch on `output_file_id` (above) instead.
    assert batch.metadata["output_file_uri"] == ""
    # The bare prefix is still preserved in metadata so callers can list it.
    assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX


def test_cancelled_status_sets_cancelled_at(patched_boto3):
    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(
        status="Stopped"
    )

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.status == "cancelled"
    assert batch.cancelled_at == int(END_TIME.timestamp())
    assert batch.completed_at is None
    assert batch.failed_at is None
    assert batch.expired_at is None


def test_expired_status_sets_expired_at(patched_boto3):
    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(
        status="Expired"
    )

    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)

    assert batch.status == "expired"
    assert batch.expired_at == int(END_TIME.timestamp())
    assert batch.completed_at is None
    assert batch.failed_at is None
    assert batch.cancelled_at is None


def test_logging_obj_pre_and_post_call_invoked(patched_boto3):
    """`pre_call` / `post_call` get called with sensible payloads when a
    `logging_obj` is supplied."""
    _, _ = patched_boto3
    logging_obj = MagicMock()

    BedrockBatchesHandler._handle_model_invocation_job_status(
        batch_id=JOB_ARN, logging_obj=logging_obj
    )

    logging_obj.pre_call.assert_called_once()
    logging_obj.post_call.assert_called_once()

    pre_kwargs = logging_obj.pre_call.call_args.kwargs
    assert pre_kwargs["input"] == JOB_ARN
    assert pre_kwargs["additional_args"]["complete_input_dict"] == {
        "jobIdentifier": JOB_ARN
    }
    # Logged URL must use the bare job id, not the full ARN, so it doesn't
    # double the `model-invocation-job/` segment or embed colons in the path.
    assert pre_kwargs["additional_args"]["api_base"] == (
        f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}"
    )

    post_kwargs = logging_obj.post_call.call_args.kwargs
    assert post_kwargs["input"] == JOB_ARN
    assert post_kwargs["original_response"]["jobArn"] == JOB_ARN


def test_missing_boto3_raises_helpful_import_error():
    """If boto3 isn't installed we should raise a clear, actionable
    ImportError rather than letting a NameError escape."""
    real_import = (
        __builtins__["__import__"]
        if isinstance(__builtins__, dict)
        else __builtins__.__import__
    )

    def fake_import(name, *args, **kwargs):
        if name == "boto3":
            raise ImportError("No module named 'boto3'")
        return real_import(name, *args, **kwargs)

    with patch("builtins.__import__", side_effect=fake_import):
        with pytest.raises(ImportError, match="pip install boto3"):
            BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)


def test_logging_url_uses_bare_id_when_only_id_passed(patched_boto3):
    """If the caller passes just the trailing job id (also valid for
    `GetModelInvocationJob`), the logged URL should use it as-is."""
    _, _ = patched_boto3
    logging_obj = MagicMock()

    BedrockBatchesHandler._handle_model_invocation_job_status(
        batch_id=JOB_ID, aws_region_name="us-west-2", logging_obj=logging_obj
    )

    pre_kwargs = logging_obj.pre_call.call_args.kwargs
    assert pre_kwargs["additional_args"]["api_base"] == (
        f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}"
    )


def test_cancel_batch_stops_job_and_returns_mapped_status(patched_boto3):
    fake_client, boto_client_factory = patched_boto3
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopping")

    batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)

    fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
    _, kwargs = boto_client_factory.call_args
    assert kwargs["region_name"] == "us-west-2"
    assert batch.status == "cancelling"


def test_cancel_batch_tolerates_already_terminal_job(patched_boto3):
    from botocore.exceptions import ClientError

    fake_client, _ = patched_boto3
    fake_client.stop_model_invocation_job.side_effect = ClientError(
        {"Error": {"Code": "ValidationException", "Message": "Job is already in a terminal state"}},
        "StopModelInvocationJob",
    )
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")

    batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)

    assert batch.status == "cancelled"


def test_cancel_batch_tolerates_conflict_on_already_stopped_job(patched_boto3):
    from botocore.exceptions import ClientError

    fake_client, _ = patched_boto3
    fake_client.stop_model_invocation_job.side_effect = ClientError(
        {"Error": {"Code": "ConflictException", "Message": "Job cannot be stopped in its current state"}},
        "StopModelInvocationJob",
    )
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")

    batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)

    assert batch.status == "cancelled"


def test_cancel_batch_reraises_conflict_when_job_not_terminal(patched_boto3):
    from botocore.exceptions import ClientError

    fake_client, _ = patched_boto3
    fake_client.stop_model_invocation_job.side_effect = ClientError(
        {"Error": {"Code": "ConflictException", "Message": "Operation conflicts with current job state"}},
        "StopModelInvocationJob",
    )
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="InProgress")

    with pytest.raises(ClientError):
        BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)


def test_cancel_batch_reraises_validation_error_when_job_not_terminal(patched_boto3):
    from botocore.exceptions import ClientError

    fake_client, _ = patched_boto3
    fake_client.stop_model_invocation_job.side_effect = ClientError(
        {"Error": {"Code": "ValidationException", "Message": "Cannot stop job in current state"}},
        "StopModelInvocationJob",
    )
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="InProgress")

    with pytest.raises(ClientError):
        BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)


def test_cancel_batch_reraises_other_client_errors(patched_boto3):
    from botocore.exceptions import ClientError

    fake_client, _ = patched_boto3
    fake_client.stop_model_invocation_job.side_effect = ClientError(
        {"Error": {"Code": "AccessDeniedException", "Message": "not authorized"}},
        "StopModelInvocationJob",
    )

    with pytest.raises(ClientError):
        BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)

    fake_client.get_model_invocation_job.assert_not_called()


def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3):
    import litellm

    fake_client, _ = patched_boto3
    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")

    batch = litellm.cancel_batch(batch_id=JOB_ARN, custom_llm_provider="bedrock")

    fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
    assert batch.status == "cancelled"


class _TagGatedSTSClient:
    """Stands in for STS behind a trust policy that only admits sessions carrying ``tags``."""

    def __init__(self, tags: list[dict[str, str]], access_key_id: str) -> None:
        self._tags = tags
        self._access_key_id = access_key_id

    def get_caller_identity(self):
        return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}

    def assume_role(self, **params):
        from botocore.exceptions import ClientError

        if list(params.get("Tags", ())) != self._tags:
            raise ClientError(
                {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}},
                "AssumeRole",
            )
        return {
            "Credentials": {
                "AccessKeyId": self._access_key_id,
                "SecretAccessKey": "assumed-secret",
                "SessionToken": "assumed-session-token",
                "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30),
            }
        }


def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_session(monkeypatch):
    """Status polling must assume the role with the deployment's session tags, like every other call."""
    monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
    monkeypatch.delenv("AWS_ROLE_ARN", raising=False)
    tags = [{"Key": "team", "Value": "genai"}]
    bedrock_client_kwargs: list[dict] = []
    fake_bedrock = MagicMock()
    fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response()

    def boto3_client(service_name, **kwargs):
        if service_name == "sts":
            return _TagGatedSTSClient(tags, "ASIABATCHSTATUSTAGGED")
        bedrock_client_kwargs.append(kwargs)
        return fake_bedrock

    with patch("boto3.client", side_effect=boto3_client):
        batch = BedrockBatchesHandler._handle_model_invocation_job_status(
            batch_id=JOB_ARN,
            aws_access_key_id="AKIABATCHSTATUSCALLER",
            aws_secret_access_key="pod-caller-secret",
            aws_role_name="arn:aws:iam::999999999999:role/litellm-batch-role",
            aws_session_name="litellm-batch-session",
            aws_session_tags=tags,
        )

    assert batch.status == "completed"
    assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHSTATUSTAGGED"]


def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatch):
    """Cancelling on a tag-gated role must forward the deployment's session tags to both the stop and status calls."""
    import litellm

    monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
    monkeypatch.delenv("AWS_ROLE_ARN", raising=False)
    tags = [{"Key": "team", "Value": "genai"}]
    bedrock_client_kwargs: list[dict] = []
    fake_bedrock = MagicMock()
    fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")

    def boto3_client(service_name, **kwargs):
        if service_name == "sts":
            return _TagGatedSTSClient(tags, "ASIABATCHCANCELTAGGED")
        bedrock_client_kwargs.append(kwargs)
        return fake_bedrock

    with patch("boto3.client", side_effect=boto3_client):
        batch = litellm.cancel_batch(
            batch_id=JOB_ARN,
            custom_llm_provider="bedrock",
            aws_access_key_id="AKIABATCHCANCELCALLER",
            aws_secret_access_key="pod-caller-secret",
            aws_role_name="arn:aws:iam::999999999999:role/litellm-batch-role",
            aws_session_name="litellm-batch-session",
            aws_session_tags=tags,
        )

    fake_bedrock.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
    assert batch.status == "cancelled"
    assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2


class _JsonBody:
    def __init__(self, payload: bytes) -> None:
        self._payload: Final = payload

    def stream(self) -> Iterator[bytes]:
        return iter((self._payload,))


class _AuthorizationRecorder:
    def __init__(self, body: Mapping[str, object]) -> None:
        self._payload: Final = json.dumps(body, default=str).encode()
        self.authorization_headers: tuple[str, ...] = ()

    def send(self, request: AWSPreparedRequest) -> AWSResponse:
        raw_authorization: Final = request.headers["Authorization"]
        authorization: Final = (
            raw_authorization.decode() if isinstance(raw_authorization, bytes) else str(raw_authorization)
        )
        self.authorization_headers = (*self.authorization_headers, authorization)
        return AWSResponse(request.url, 200, {"content-type": "application/json"}, _JsonBody(self._payload))


def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch):
    """A proxy-wide AWS_BEARER_TOKEN_BEDROCK must not override the deployment's own SigV4 credentials."""
    monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token")
    recorder: Final = _AuthorizationRecorder(_fake_boto3_response())

    with patch("botocore.httpsession.URLLib3Session.send", recorder.send):
        batch = BedrockBatchesHandler._handle_model_invocation_job_status(
            batch_id=JOB_ARN,
            aws_access_key_id="AKIADEPLOYMENTKEY",
            aws_secret_access_key="deployment-secret",
        )

    assert batch.status == "completed"
    assert len(recorder.authorization_headers) == 1
    assert recorder.authorization_headers[0].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/")


def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch):
    monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token")
    recorder: Final = _AuthorizationRecorder(_fake_boto3_response(status="Stopped"))

    with patch("botocore.httpsession.URLLib3Session.send", recorder.send):
        batch = BedrockBatchesHandler.cancel_batch(
            batch_id=JOB_ARN,
            aws_access_key_id="AKIADEPLOYMENTKEY",
            aws_secret_access_key="deployment-secret",
        )

    assert batch.status == "cancelled"
    assert len(recorder.authorization_headers) == 2
    assert all(h.startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in recorder.authorization_headers)