datasetio

This commit is contained in:
Xi Yan 2025-03-12 16:47:40 -07:00
parent 97d6b87e05
commit 9dcf617556
2 changed files with 7 additions and 3 deletions

View file

@ -15,11 +15,12 @@ import pytest
#
# LLAMA_STACK_CONFIG="template-name" pytest -v tests/integration/datasetio
@pytest.fixture
def test_dataset(llama_stack_client):
register_dataset(llama_stack_client)
yield # This is where the test function will run
# Teardown - this always runs, even if the test fails
try:
llama_stack_client.datasets.unregister("test_dataset")

View file

@ -9,11 +9,12 @@ import pytest
from ..datasetio.test_datasetio import register_dataset
@pytest.fixture
def test_dataset_rag(llama_stack_client):
register_dataset(llama_stack_client, for_rag=True)
yield # This is where the test function will run
# Teardown - this always runs, even if the test fails
try:
llama_stack_client.datasets.unregister("test_dataset")
@ -124,7 +125,9 @@ def test_scoring_score(llama_stack_client, test_dataset_rag):
assert len(response.results[x].score_rows) == 5
def test_scoring_score_with_params_llm_as_judge(llama_stack_client, sample_judge_prompt_template, judge_model_id, test_dataset_rag):
def test_scoring_score_with_params_llm_as_judge(
llama_stack_client, sample_judge_prompt_template, judge_model_id, test_dataset_rag
):
# scoring individual rows
rows = llama_stack_client.datasetio.get_rows_paginated(
dataset_id="test_dataset",