diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 32063f65765..fd8798cbdd8 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10017,6 +10017,26 @@ def test_is_base64_encoded_2(): assert is_base64_encoded(s="Dog") is False +@pytest.mark.parametrize( + "value", + [ + "", + "data:text/plain;base64", + "data:text/plain;base64,!!!not-base64!!!", + ], +) +def test_is_base64_encoded_rejects_malformed_input(value): + from litellm.utils import is_base64_encoded + + assert is_base64_encoded(s=value) is False + + +def test_is_base64_encoded_accepts_text_data_uri(): + from litellm.utils import is_base64_encoded + + assert is_base64_encoded(s="data:text/plain;base64,aGVsbG8gd29ybGQ=") is True + + @pytest.mark.parametrize( "messages, expected_bool", [