async def test_async_retry_on_connection_error(mock_client_class):
    """Test that the async embedding method retries on connection errors."""
    # Setup mock for async client
    mock_client = mock_client_class.return_value
    mock_aio = MagicMock()
    mock_client.aio = mock_aio
    mock_aio_models = mock_aio.models

    # Create mock embedding result for successful attempt
    mock_embedding = MagicMock()
    mock_embedding.values = [0.4, 0.5, 0.6]
    mock_result = MagicMock()
    mock_result.embeddings = [mock_embedding]

    # Create two different AsyncMock objects
    fail_mock = AsyncMock(
        side_effect=requests.exceptions.ConnectionError("Connection error")
    )
    success_mock = AsyncMock(return_value=mock_result)

    # Configure the mock to return different mocks on consecutive calls
    mock_aio_models.embed_content = fail_mock

    # Test async embedding with retries
    emb = GoogleGenAIEmbedding(
        api_key="fake_key",
        retries=2,
        retry_min_seconds=0.1,  # Use small values for faster tests
        retry_max_seconds=0.2,
    )

    # Replace the mock after the first call to simulate recovery
    async def side_effect(*args, **kwargs):
        # Replace the mock after first call
        mock_aio_models.embed_content = success_mock
        raise requests.exceptions.ConnectionError("Connection error")

    fail_mock.side_effect = side_effect

    # This should fail once, retry, then succeed
    result = await emb.aget_query_embedding("test query")

    # Verify the result is correct
    assert result == [0.4, 0.5, 0.6]

    # Verify both mocks were called (original + retry)
    fail_mock.assert_called_once()
    success_mock.assert_called_once()
Test that the async embedding method retries on connection errors.
llama-index-integrations/embeddings/llama-index-embeddings-google-genai/tests/test_embeddings_gemini.py