def test_batch_size_validation():
    """Test that batch size validation works correctly."""
    # Test batch size exceeding the limit
    with pytest.raises(ValueError) as exc_info:
        CohereEmbedding(api_key=os.environ["CO_API_KEY"], embed_batch_size=97)
    assert "exceeds the maximum allowed value of 96" in str(exc_info.value)

    # Test batch size at the limit (should not raise)
    emb = CohereEmbedding(api_key=os.environ["CO_API_KEY"], embed_batch_size=96)
    assert emb.embed_batch_size == 96

    # Test batch size below the limit (should not raise)
    emb = CohereEmbedding(api_key=os.environ["CO_API_KEY"], embed_batch_size=50)
    assert emb.embed_batch_size == 50
Test that batch size validation works correctly.
llama-index-integrations/embeddings/llama-index-embeddings-cohere/tests/test_embeddings.py