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

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

@ -5,103 +5,110 @@
# 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():
os.environ["NVIDIA_DATASETS_URL"] = "http://nemo.test/datasets" """Fixture to set up NvidiaDatasetIOAdapter with mocked requests."""
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(
"llama_stack.providers.remote.datasetio.nvidia.datasetio.NvidiaDatasetIOAdapter._make_request"
)
self.mock_make_request = self.make_request_patcher.start()
def tearDown(self): with patch(
self.make_request_patcher.stop() "llama_stack.providers.remote.datasetio.nvidia.datasetio.NvidiaDatasetIOAdapter._make_request"
) as mock_make_request:
yield adapter, mock_make_request
@pytest.fixture(autouse=True)
def inject_fixtures(self, run_async):
self.run_async = run_async
def _assert_request(self, mock_call, expected_method, expected_path, expected_json=None): def _assert_request(mock_call, expected_method, expected_path, expected_json=None):
"""Helper method to verify request details in mock calls.""" """Helper function 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
assert call_args[0][1] == expected_path assert call_args[0][1] == expected_path
if expected_json: if expected_json:
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):
"id": "dataset-123456", adapter, mock_make_request = nvidia_adapter
mock_make_request.return_value = {
"id": "dataset-123456",
"name": "test-dataset",
"namespace": "default",
}
dataset_def = Dataset(
identifier="test-dataset",
type=ResourceType.dataset,
provider_resource_id="",
provider_id="",
purpose=DatasetPurpose.post_training_messages,
source=URIDataSource(uri="https://example.com/data.jsonl"),
metadata={"provider_id": "nvidia", "format": "jsonl", "description": "Test dataset description"},
)
run_async(adapter.register_dataset(dataset_def))
mock_make_request.assert_called_once()
_assert_request(
mock_make_request,
"POST",
"/v1/datasets",
expected_json={
"name": "test-dataset", "name": "test-dataset",
"namespace": "default", "namespace": "default",
} "files_url": "https://example.com/data.jsonl",
"project": "default",
"format": "jsonl",
"description": "Test dataset description",
},
)
dataset_def = Dataset(
identifier="test-dataset",
type="dataset",
provider_resource_id="",
provider_id="",
purpose=DatasetPurpose.post_training_messages,
source=URIDataSource(uri="https://example.com/data.jsonl"),
metadata={"provider_id": "nvidia", "format": "jsonl", "description": "Test dataset description"},
)
self.run_async(self.adapter.register_dataset(dataset_def)) def test_unregister_dataset(nvidia_adapter, run_async):
adapter, mock_make_request = nvidia_adapter
mock_make_request.return_value = {
"message": "Resource deleted successfully.",
"id": "dataset-81RSQp7FKX3rdBtKvF9Skn",
"deleted_at": None,
}
dataset_id = "test-dataset"
self.mock_make_request.assert_called_once() run_async(adapter.unregister_dataset(dataset_id))
self._assert_request(
self.mock_make_request,
"POST",
"/v1/datasets",
expected_json={
"name": "test-dataset",
"namespace": "default",
"files_url": "https://example.com/data.jsonl",
"project": "default",
"format": "jsonl",
"description": "Test dataset description",
},
)
def test_unregister_dataset(self): mock_make_request.assert_called_once()
self.mock_make_request.return_value = { _assert_request(mock_make_request, "DELETE", "/v1/datasets/default/test-dataset")
"message": "Resource deleted successfully.",
"id": "dataset-81RSQp7FKX3rdBtKvF9Skn",
"deleted_at": None,
}
dataset_id = "test-dataset"
self.run_async(self.adapter.unregister_dataset(dataset_id))
self.mock_make_request.assert_called_once() def test_register_dataset_with_custom_namespace_project(run_async):
self._assert_request(self.mock_make_request, "DELETE", "/v1/datasets/default/test-dataset") """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", project_id="custom-project",
project_id="custom-project", )
) 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()