diff --git a/llama_stack/providers/tests/datasetio/test_datasetio.py b/llama_stack/providers/tests/datasetio/test_datasetio.py
index 2b2d57ddd..dd2cbd019 100644
--- a/llama_stack/providers/tests/datasetio/test_datasetio.py
+++ b/llama_stack/providers/tests/datasetio/test_datasetio.py
@@ -57,7 +57,7 @@ async def register_dataset(
 
     await datasets_impl.register_dataset(
         dataset_id=dataset_id,
-        schema=dataset_schema,
+        dataset_schema=dataset_schema,
         url=URL(uri=test_url),
     )
 
diff --git a/llama_stack/providers/tests/eval/test_eval.py b/llama_stack/providers/tests/eval/test_eval.py
index 2d08aabe7..168745550 100644
--- a/llama_stack/providers/tests/eval/test_eval.py
+++ b/llama_stack/providers/tests/eval/test_eval.py
@@ -163,7 +163,7 @@ class Testeval:
 
         await datasets_impl.register_dataset(
             dataset_id="mmlu",
-            schema={
+            dataset_schema={
                 "input_query": StringType(),
                 "expected_answer": StringType(),
                 "chat_completion_input": ChatCompletionInputType(),