From f31476541e61fc262e35853072161503a779ece8 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Wed, 13 Nov 2024 11:13:32 -0500
Subject: [PATCH] tests fix

---
 llama_stack/providers/tests/datasetio/test_datasetio.py | 2 +-
 llama_stack/providers/tests/eval/test_eval.py           | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

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(),