phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-20 11:47:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

chore(test): migrate unit tests from unittest to pytest for nvidia datastore (#2790)

Browse source 
This PR replaces unittest with pytest.

Part of https://github.com/meta-llama/llama-stack/issues/2680

cc @leseb

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>
This commit is contained in:
Mustafa Elbehery 2025-07-18 11:52:47 +02:00 • committed by GitHub
parent 55713abe7d
commit 3cdf748a8e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 83 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

79
tests/unit/providers/nvidia/test_datastore.py
View file

@ -5,38 +5,34 @@
# the root directory of this source tree. # the root directory of this source tree.
import os import os
import unittest
from unittest.mock import patch from unittest.mock import patch
import pytest import pytest
from llama_stack.apis.datasets import Dataset, DatasetPurpose, URIDataSource from llama_stack.apis.datasets import Dataset, DatasetPurpose, URIDataSource
from llama_stack.apis.resource import ResourceType
from llama_stack.providers.remote.datasetio.nvidia.config import NvidiaDatasetIOConfig from llama_stack.providers.remote.datasetio.nvidia.config import NvidiaDatasetIOConfig
from llama_stack.providers.remote.datasetio.nvidia.datasetio import NvidiaDatasetIOAdapter from llama_stack.providers.remote.datasetio.nvidia.datasetio import NvidiaDatasetIOAdapter
class TestNvidiaDatastore(unittest.TestCase): @pytest.fixture
def setUp(self): def nvidia_adapter():
"""Fixture to set up NvidiaDatasetIOAdapter with mocked requests."""
os.environ["NVIDIA_DATASETS_URL"] = "http://nemo.test/datasets" os.environ["NVIDIA_DATASETS_URL"] = "http://nemo.test/datasets"
config = NvidiaDatasetIOConfig( config = NvidiaDatasetIOConfig(
datasets_url=os.environ["NVIDIA_DATASETS_URL"], dataset_namespace="default", project_id="default" datasets_url=os.environ["NVIDIA_DATASETS_URL"], dataset_namespace="default", project_id="default"
) )
self.adapter = NvidiaDatasetIOAdapter(config) adapter = NvidiaDatasetIOAdapter(config)
self.make_request_patcher = patch(
with patch(
"llama_stack.providers.remote.datasetio.nvidia.datasetio.NvidiaDatasetIOAdapter._make_request" "llama_stack.providers.remote.datasetio.nvidia.datasetio.NvidiaDatasetIOAdapter._make_request"
) ) as mock_make_request:
self.mock_make_request = self.make_request_patcher.start() yield adapter, mock_make_request
def tearDown(self):
self.make_request_patcher.stop()
@pytest.fixture(autouse=True) def _assert_request(mock_call, expected_method, expected_path, expected_json=None):
def inject_fixtures(self, run_async): """Helper function to verify request details in mock calls."""
self.run_async = run_async
def _assert_request(self, mock_call, expected_method, expected_path, expected_json=None):
"""Helper method to verify request details in mock calls."""
call_args = mock_call.call_args call_args = mock_call.call_args
assert call_args[0][0] == expected_method assert call_args[0][0] == expected_method
@ -46,8 +42,10 @@ class TestNvidiaDatastore(unittest.TestCase):
for key, value in expected_json.items(): for key, value in expected_json.items():
assert call_args[1]["json"][key] == value assert call_args[1]["json"][key] == value
def test_register_dataset(self):
self.mock_make_request.return_value = { def test_register_dataset(nvidia_adapter, run_async):
adapter, mock_make_request = nvidia_adapter
mock_make_request.return_value = {
"id": "dataset-123456", "id": "dataset-123456",
"name": "test-dataset", "name": "test-dataset",
"namespace": "default", "namespace": "default",
@ -55,7 +53,7 @@ class TestNvidiaDatastore(unittest.TestCase):
dataset_def = Dataset( dataset_def = Dataset(
identifier="test-dataset", identifier="test-dataset",
type="dataset", type=ResourceType.dataset,
provider_resource_id="", provider_resource_id="",
provider_id="", provider_id="",
purpose=DatasetPurpose.post_training_messages, purpose=DatasetPurpose.post_training_messages,
@ -63,11 +61,11 @@ class TestNvidiaDatastore(unittest.TestCase):
metadata={"provider_id": "nvidia", "format": "jsonl", "description": "Test dataset description"}, metadata={"provider_id": "nvidia", "format": "jsonl", "description": "Test dataset description"},
) )
self.run_async(self.adapter.register_dataset(dataset_def)) run_async(adapter.register_dataset(dataset_def))
self.mock_make_request.assert_called_once() mock_make_request.assert_called_once()
self._assert_request( _assert_request(
self.mock_make_request, mock_make_request,
"POST", "POST",
"/v1/datasets", "/v1/datasets",
expected_json={ expected_json={
@ -80,20 +78,26 @@ class TestNvidiaDatastore(unittest.TestCase):
}, },
) )
def test_unregister_dataset(self):
self.mock_make_request.return_value = { def test_unregister_dataset(nvidia_adapter, run_async):
adapter, mock_make_request = nvidia_adapter
mock_make_request.return_value = {
"message": "Resource deleted successfully.", "message": "Resource deleted successfully.",
"id": "dataset-81RSQp7FKX3rdBtKvF9Skn", "id": "dataset-81RSQp7FKX3rdBtKvF9Skn",
"deleted_at": None, "deleted_at": None,
} }
dataset_id = "test-dataset" dataset_id = "test-dataset"
self.run_async(self.adapter.unregister_dataset(dataset_id)) run_async(adapter.unregister_dataset(dataset_id))
self.mock_make_request.assert_called_once() mock_make_request.assert_called_once()
self._assert_request(self.mock_make_request, "DELETE", "/v1/datasets/default/test-dataset") _assert_request(mock_make_request, "DELETE", "/v1/datasets/default/test-dataset")
def test_register_dataset_with_custom_namespace_project(run_async):
"""Test with custom namespace and project configuration."""
os.environ["NVIDIA_DATASETS_URL"] = "http://nemo.test/datasets"
def test_register_dataset_with_custom_namespace_project(self):
custom_config = NvidiaDatasetIOConfig( custom_config = NvidiaDatasetIOConfig(
datasets_url=os.environ["NVIDIA_DATASETS_URL"], datasets_url=os.environ["NVIDIA_DATASETS_URL"],
dataset_namespace="custom-namespace", dataset_namespace="custom-namespace",
@ -101,7 +105,10 @@ class TestNvidiaDatastore(unittest.TestCase):
) )
custom_adapter = NvidiaDatasetIOAdapter(custom_config) custom_adapter = NvidiaDatasetIOAdapter(custom_config)
self.mock_make_request.return_value = { with patch(
"llama_stack.providers.remote.datasetio.nvidia.datasetio.NvidiaDatasetIOAdapter._make_request"
) as mock_make_request:
mock_make_request.return_value = {
"id": "dataset-123456", "id": "dataset-123456",
"name": "test-dataset", "name": "test-dataset",
"namespace": "custom-namespace", "namespace": "custom-namespace",
@ -109,7 +116,7 @@ class TestNvidiaDatastore(unittest.TestCase):
dataset_def = Dataset( dataset_def = Dataset(
identifier="test-dataset", identifier="test-dataset",
type="dataset", type=ResourceType.dataset,
provider_resource_id="", provider_resource_id="",
provider_id="", provider_id="",
purpose=DatasetPurpose.post_training_messages, purpose=DatasetPurpose.post_training_messages,
@ -117,11 +124,11 @@ class TestNvidiaDatastore(unittest.TestCase):
metadata={"format": "jsonl"}, metadata={"format": "jsonl"},
) )
self.run_async(custom_adapter.register_dataset(dataset_def)) run_async(custom_adapter.register_dataset(dataset_def))
self.mock_make_request.assert_called_once() mock_make_request.assert_called_once()
self._assert_request( _assert_request(
self.mock_make_request, mock_make_request,
"POST", "POST",
"/v1/datasets", "/v1/datasets",
expected_json={ expected_json={
@ -132,7 +139,3 @@ class TestNvidiaDatastore(unittest.TestCase):
"format": "jsonl", "format": "jsonl",
}, },
) )
if __name__ == "__main__":
unittest.main()