diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5cf872f4ea7..7c54ef309cf 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -9872,6 +9872,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", [