diff --git a/sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py b/sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py
index 3d2a253dc0e..a29383a5c9c 100644
--- a/sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py
+++ b/sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py
@@ -38,7 +38,7 @@ def test_online_retrieval(environment, universal_data_sources, benchmark):
     sample_customers = random.sample(entities.customer_vals, 10)
 
     entity_rows = [
-        {"driver": d, "customer_id": c, "val_to_add": 50}
+        {"driver_id": d, "customer_id": c, "val_to_add": 50}
         for (d, c) in zip(sample_drivers, sample_customers)
     ]