From 6ab5760a1bdac70bc85803b3a20f4291b30b41a6 Mon Sep 17 00:00:00 2001 From: Mustafa Elbehery Date: Thu, 24 Jul 2025 18:41:07 +0200 Subject: [PATCH] chore(test): migrate unit tests from unittest to pytest nvidia test safety (#2793) This PR replaces unittest with pytest. Part of https://github.com/meta-llama/llama-stack/issues/2680 cc @leseb Signed-off-by: Mustafa Elbehery --- tests/unit/providers/nvidia/test_safety.py | 554 +++++++++++---------- 1 file changed, 293 insertions(+), 261 deletions(-) diff --git a/tests/unit/providers/nvidia/test_safety.py b/tests/unit/providers/nvidia/test_safety.py index 73fc32a02..bfd91f466 100644 --- a/tests/unit/providers/nvidia/test_safety.py +++ b/tests/unit/providers/nvidia/test_safety.py @@ -5,321 +5,353 @@ # the root directory of this source tree. import os -import unittest from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from llama_stack.apis.inference import CompletionMessage, UserMessage +from llama_stack.apis.resource import ResourceType from llama_stack.apis.safety import RunShieldResponse, ViolationLevel from llama_stack.apis.shields import Shield +from llama_stack.models.llama.datatypes import StopReason from llama_stack.providers.remote.safety.nvidia.config import NVIDIASafetyConfig from llama_stack.providers.remote.safety.nvidia.nvidia import NVIDIASafetyAdapter -class TestNVIDIASafetyAdapter(unittest.TestCase): - def setUp(self): - os.environ["NVIDIA_GUARDRAILS_URL"] = "http://nemo.test" +class TestNVIDIASafetyAdapter(NVIDIASafetyAdapter): + """Test implementation that provides the required shield_store.""" - # Initialize the adapter - self.config = NVIDIASafetyConfig( - guardrails_service_url=os.environ["NVIDIA_GUARDRAILS_URL"], - ) - self.adapter = NVIDIASafetyAdapter(config=self.config) - self.shield_store = AsyncMock() - self.adapter.shield_store = self.shield_store + def __init__(self, config: NVIDIASafetyConfig, shield_store): + super().__init__(config) + self.shield_store = shield_store - # Mock the HTTP request methods - self.guardrails_post_patcher = patch( - "llama_stack.providers.remote.safety.nvidia.nvidia.NeMoGuardrails._guardrails_post" - ) - self.mock_guardrails_post = self.guardrails_post_patcher.start() - self.mock_guardrails_post.return_value = {"status": "allowed"} - def tearDown(self): - """Clean up after each test.""" - self.guardrails_post_patcher.stop() +@pytest.fixture +def nvidia_adapter(): + """Set up the NVIDIASafetyAdapter for testing.""" + os.environ["NVIDIA_GUARDRAILS_URL"] = "http://nemo.test" - @pytest.fixture(autouse=True) - def inject_fixtures(self, run_async): - self.run_async = run_async + # Initialize the adapter + config = NVIDIASafetyConfig( + guardrails_service_url=os.environ["NVIDIA_GUARDRAILS_URL"], + ) - def _assert_request( - self, - mock_call: MagicMock, - expected_url: str, - expected_headers: dict[str, str] | None = None, - expected_json: dict[str, Any] | None = None, - ) -> None: - """ - Helper method to verify request details in mock API calls. + # Create a mock shield store that implements the ShieldStore protocol + shield_store = AsyncMock() + shield_store.get_shield = AsyncMock() - Args: - mock_call: The MagicMock object that was called - expected_url: The expected URL to which the request was made - expected_headers: Optional dictionary of expected request headers - expected_json: Optional dictionary of expected JSON payload - """ - call_args = mock_call.call_args + adapter = TestNVIDIASafetyAdapter(config=config, shield_store=shield_store) - # Check URL - assert call_args[0][0] == expected_url + return adapter - # Check headers if provided - if expected_headers: - for key, value in expected_headers.items(): - assert call_args[1]["headers"][key] == value - # Check JSON if provided - if expected_json: - for key, value in expected_json.items(): - if isinstance(value, dict): - for nested_key, nested_value in value.items(): - assert call_args[1]["json"][key][nested_key] == nested_value - else: - assert call_args[1]["json"][key] == value +@pytest.fixture +def mock_guardrails_post(): + """Mock the HTTP request methods.""" + with patch("llama_stack.providers.remote.safety.nvidia.nvidia.NeMoGuardrails._guardrails_post") as mock_post: + mock_post.return_value = {"status": "allowed"} + yield mock_post - def test_register_shield_with_valid_id(self): - shield = Shield( - provider_id="nvidia", - type="shield", - identifier="test-shield", - provider_resource_id="test-model", - ) - # Register the shield - self.run_async(self.adapter.register_shield(shield)) +def _assert_request( + mock_call: MagicMock, + expected_url: str, + expected_headers: dict[str, str] | None = None, + expected_json: dict[str, Any] | None = None, +) -> None: + """ + Helper method to verify request details in mock API calls. - def test_register_shield_without_id(self): - shield = Shield( - provider_id="nvidia", - type="shield", - identifier="test-shield", - provider_resource_id="", - ) + Args: + mock_call: The MagicMock object that was called + expected_url: The expected URL to which the request was made + expected_headers: Optional dictionary of expected request headers + expected_json: Optional dictionary of expected JSON payload + """ + call_args = mock_call.call_args - # Register the shield should raise a ValueError - with self.assertRaises(ValueError): - self.run_async(self.adapter.register_shield(shield)) + # Check URL + assert call_args[0][0] == expected_url - def test_run_shield_allowed(self): - # Set up the shield - shield_id = "test-shield" - shield = Shield( - provider_id="nvidia", - type="shield", - identifier=shield_id, - provider_resource_id="test-model", - ) - self.shield_store.get_shield.return_value = shield + # Check headers if provided + if expected_headers: + for key, value in expected_headers.items(): + assert call_args[1]["headers"][key] == value - # Mock Guardrails API response - self.mock_guardrails_post.return_value = {"status": "allowed"} + # Check JSON if provided + if expected_json: + for key, value in expected_json.items(): + if isinstance(value, dict): + for nested_key, nested_value in value.items(): + assert call_args[1]["json"][key][nested_key] == nested_value + else: + assert call_args[1]["json"][key] == value - # Run the shield - messages = [ - UserMessage(role="user", content="Hello, how are you?"), - CompletionMessage( - role="assistant", - content="I'm doing well, thank you for asking!", - stop_reason="end_of_message", - tool_calls=[], - ), - ] - result = self.run_async(self.adapter.run_shield(shield_id, messages)) - # Verify the shield store was called - self.shield_store.get_shield.assert_called_once_with(shield_id) +async def test_register_shield_with_valid_id(nvidia_adapter): + adapter = nvidia_adapter - # Verify the Guardrails API was called correctly - self.mock_guardrails_post.assert_called_once_with( - path="/v1/guardrail/checks", - data={ - "model": shield_id, - "messages": [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, - ], - "temperature": 1.0, - "top_p": 1, - "frequency_penalty": 0, - "presence_penalty": 0, - "max_tokens": 160, - "stream": False, - "guardrails": { - "config_id": "self-check", - }, + shield = Shield( + provider_id="nvidia", + type=ResourceType.shield, + identifier="test-shield", + provider_resource_id="test-model", + ) + + # Register the shield + await adapter.register_shield(shield) + + +async def test_register_shield_without_id(nvidia_adapter): + adapter = nvidia_adapter + + shield = Shield( + provider_id="nvidia", + type=ResourceType.shield, + identifier="test-shield", + provider_resource_id="", + ) + + # Register the shield should raise a ValueError + with pytest.raises(ValueError): + await adapter.register_shield(shield) + + +async def test_run_shield_allowed(nvidia_adapter, mock_guardrails_post): + adapter = nvidia_adapter + + # Set up the shield + shield_id = "test-shield" + shield = Shield( + provider_id="nvidia", + type=ResourceType.shield, + identifier=shield_id, + provider_resource_id="test-model", + ) + adapter.shield_store.get_shield.return_value = shield + + # Mock Guardrails API response + mock_guardrails_post.return_value = {"status": "allowed"} + + # Run the shield + messages = [ + UserMessage(role="user", content="Hello, how are you?"), + CompletionMessage( + role="assistant", + content="I'm doing well, thank you for asking!", + stop_reason=StopReason.end_of_message, + tool_calls=[], + ), + ] + result = await adapter.run_shield(shield_id, messages) + + # Verify the shield store was called + adapter.shield_store.get_shield.assert_called_once_with(shield_id) + + # Verify the Guardrails API was called correctly + mock_guardrails_post.assert_called_once_with( + path="/v1/guardrail/checks", + data={ + "model": shield_id, + "messages": [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, + ], + "temperature": 1.0, + "top_p": 1, + "frequency_penalty": 0, + "presence_penalty": 0, + "max_tokens": 160, + "stream": False, + "guardrails": { + "config_id": "self-check", }, - ) + }, + ) - # Verify the result - assert isinstance(result, RunShieldResponse) - assert result.violation is None + # Verify the result + assert isinstance(result, RunShieldResponse) + assert result.violation is None - def test_run_shield_blocked(self): - # Set up the shield - shield_id = "test-shield" - shield = Shield( - provider_id="nvidia", - type="shield", - identifier=shield_id, - provider_resource_id="test-model", - ) - self.shield_store.get_shield.return_value = shield - # Mock Guardrails API response - self.mock_guardrails_post.return_value = {"status": "blocked", "rails_status": {"reason": "harmful_content"}} +async def test_run_shield_blocked(nvidia_adapter, mock_guardrails_post): + adapter = nvidia_adapter - # Run the shield - messages = [ - UserMessage(role="user", content="Hello, how are you?"), - CompletionMessage( - role="assistant", - content="I'm doing well, thank you for asking!", - stop_reason="end_of_message", - tool_calls=[], - ), - ] - result = self.run_async(self.adapter.run_shield(shield_id, messages)) + # Set up the shield + shield_id = "test-shield" + shield = Shield( + provider_id="nvidia", + type=ResourceType.shield, + identifier=shield_id, + provider_resource_id="test-model", + ) + adapter.shield_store.get_shield.return_value = shield - # Verify the shield store was called - self.shield_store.get_shield.assert_called_once_with(shield_id) + # Mock Guardrails API response + mock_guardrails_post.return_value = {"status": "blocked", "rails_status": {"reason": "harmful_content"}} - # Verify the Guardrails API was called correctly - self.mock_guardrails_post.assert_called_once_with( - path="/v1/guardrail/checks", - data={ - "model": shield_id, - "messages": [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, - ], - "temperature": 1.0, - "top_p": 1, - "frequency_penalty": 0, - "presence_penalty": 0, - "max_tokens": 160, - "stream": False, - "guardrails": { - "config_id": "self-check", - }, + # Run the shield + messages = [ + UserMessage(role="user", content="Hello, how are you?"), + CompletionMessage( + role="assistant", + content="I'm doing well, thank you for asking!", + stop_reason=StopReason.end_of_message, + tool_calls=[], + ), + ] + result = await adapter.run_shield(shield_id, messages) + + # Verify the shield store was called + adapter.shield_store.get_shield.assert_called_once_with(shield_id) + + # Verify the Guardrails API was called correctly + mock_guardrails_post.assert_called_once_with( + path="/v1/guardrail/checks", + data={ + "model": shield_id, + "messages": [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, + ], + "temperature": 1.0, + "top_p": 1, + "frequency_penalty": 0, + "presence_penalty": 0, + "max_tokens": 160, + "stream": False, + "guardrails": { + "config_id": "self-check", }, - ) + }, + ) - # Verify the result - assert result.violation is not None - assert isinstance(result, RunShieldResponse) - assert result.violation.user_message == "Sorry I cannot do this." - assert result.violation.violation_level == ViolationLevel.ERROR - assert result.violation.metadata == {"reason": "harmful_content"} + # Verify the result + assert result.violation is not None + assert isinstance(result, RunShieldResponse) + assert result.violation.user_message == "Sorry I cannot do this." + assert result.violation.violation_level == ViolationLevel.ERROR + assert result.violation.metadata == {"reason": "harmful_content"} - def test_run_shield_not_found(self): - # Set up shield store to return None - shield_id = "non-existent-shield" - self.shield_store.get_shield.return_value = None - messages = [ - UserMessage(role="user", content="Hello, how are you?"), - ] +async def test_run_shield_not_found(nvidia_adapter, mock_guardrails_post): + adapter = nvidia_adapter - with self.assertRaises(ValueError): - self.run_async(self.adapter.run_shield(shield_id, messages)) + # Set up shield store to return None + shield_id = "non-existent-shield" + adapter.shield_store.get_shield.return_value = None - # Verify the shield store was called - self.shield_store.get_shield.assert_called_once_with(shield_id) + messages = [ + UserMessage(role="user", content="Hello, how are you?"), + ] - # Verify the Guardrails API was not called - self.mock_guardrails_post.assert_not_called() + with pytest.raises(ValueError): + await adapter.run_shield(shield_id, messages) - def test_run_shield_http_error(self): - shield_id = "test-shield" - shield = Shield( - provider_id="nvidia", - type="shield", - identifier=shield_id, - provider_resource_id="test-model", - ) - self.shield_store.get_shield.return_value = shield + # Verify the shield store was called + adapter.shield_store.get_shield.assert_called_once_with(shield_id) - # Mock Guardrails API to raise an exception - error_msg = "API Error: 500 Internal Server Error" - self.mock_guardrails_post.side_effect = Exception(error_msg) + # Verify the Guardrails API was not called + mock_guardrails_post.assert_not_called() - # Running the shield should raise an exception - messages = [ - UserMessage(role="user", content="Hello, how are you?"), - CompletionMessage( - role="assistant", - content="I'm doing well, thank you for asking!", - stop_reason="end_of_message", - tool_calls=[], - ), - ] - with self.assertRaises(Exception) as context: - self.run_async(self.adapter.run_shield(shield_id, messages)) - # Verify the shield store was called - self.shield_store.get_shield.assert_called_once_with(shield_id) +async def test_run_shield_http_error(nvidia_adapter, mock_guardrails_post): + adapter = nvidia_adapter - # Verify the Guardrails API was called correctly - self.mock_guardrails_post.assert_called_once_with( - path="/v1/guardrail/checks", - data={ - "model": shield_id, - "messages": [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, - ], - "temperature": 1.0, - "top_p": 1, - "frequency_penalty": 0, - "presence_penalty": 0, - "max_tokens": 160, - "stream": False, - "guardrails": { - "config_id": "self-check", - }, + shield_id = "test-shield" + shield = Shield( + provider_id="nvidia", + type=ResourceType.shield, + identifier=shield_id, + provider_resource_id="test-model", + ) + adapter.shield_store.get_shield.return_value = shield + + # Mock Guardrails API to raise an exception + error_msg = "API Error: 500 Internal Server Error" + mock_guardrails_post.side_effect = Exception(error_msg) + + # Running the shield should raise an exception + messages = [ + UserMessage(role="user", content="Hello, how are you?"), + CompletionMessage( + role="assistant", + content="I'm doing well, thank you for asking!", + stop_reason=StopReason.end_of_message, + tool_calls=[], + ), + ] + with pytest.raises(Exception) as exc_info: + await adapter.run_shield(shield_id, messages) + + # Verify the shield store was called + adapter.shield_store.get_shield.assert_called_once_with(shield_id) + + # Verify the Guardrails API was called correctly + mock_guardrails_post.assert_called_once_with( + path="/v1/guardrail/checks", + data={ + "model": shield_id, + "messages": [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, + ], + "temperature": 1.0, + "top_p": 1, + "frequency_penalty": 0, + "presence_penalty": 0, + "max_tokens": 160, + "stream": False, + "guardrails": { + "config_id": "self-check", }, - ) - # Verify the exception message - assert error_msg in str(context.exception) + }, + ) + # Verify the exception message + assert error_msg in str(exc_info.value) - def test_init_nemo_guardrails(self): - from llama_stack.providers.remote.safety.nvidia.nvidia import NeMoGuardrails - test_config_id = "test-custom-config-id" - config = NVIDIASafetyConfig( - guardrails_service_url=os.environ["NVIDIA_GUARDRAILS_URL"], - config_id=test_config_id, - ) - # Initialize with default parameters - test_model = "test-model" - guardrails = NeMoGuardrails(config, test_model) +def test_init_nemo_guardrails(): + from llama_stack.providers.remote.safety.nvidia.nvidia import NeMoGuardrails - # Verify the attributes are set correctly - assert guardrails.config_id == test_config_id - assert guardrails.model == test_model - assert guardrails.threshold == 0.9 # Default value - assert guardrails.temperature == 1.0 # Default value - assert guardrails.guardrails_service_url == os.environ["NVIDIA_GUARDRAILS_URL"] + os.environ["NVIDIA_GUARDRAILS_URL"] = "http://nemo.test" - # Initialize with custom parameters - guardrails = NeMoGuardrails(config, test_model, threshold=0.8, temperature=0.7) + test_config_id = "test-custom-config-id" + config = NVIDIASafetyConfig( + guardrails_service_url=os.environ["NVIDIA_GUARDRAILS_URL"], + config_id=test_config_id, + ) + # Initialize with default parameters + test_model = "test-model" + guardrails = NeMoGuardrails(config, test_model) - # Verify the attributes are set correctly - assert guardrails.config_id == test_config_id - assert guardrails.model == test_model - assert guardrails.threshold == 0.8 - assert guardrails.temperature == 0.7 - assert guardrails.guardrails_service_url == os.environ["NVIDIA_GUARDRAILS_URL"] + # Verify the attributes are set correctly + assert guardrails.config_id == test_config_id + assert guardrails.model == test_model + assert guardrails.threshold == 0.9 # Default value + assert guardrails.temperature == 1.0 # Default value + assert guardrails.guardrails_service_url == os.environ["NVIDIA_GUARDRAILS_URL"] - def test_init_nemo_guardrails_invalid_temperature(self): - from llama_stack.providers.remote.safety.nvidia.nvidia import NeMoGuardrails + # Initialize with custom parameters + guardrails = NeMoGuardrails(config, test_model, threshold=0.8, temperature=0.7) - config = NVIDIASafetyConfig( - guardrails_service_url=os.environ["NVIDIA_GUARDRAILS_URL"], - config_id="test-custom-config-id", - ) - with self.assertRaises(ValueError): - NeMoGuardrails(config, "test-model", temperature=0) + # Verify the attributes are set correctly + assert guardrails.config_id == test_config_id + assert guardrails.model == test_model + assert guardrails.threshold == 0.8 + assert guardrails.temperature == 0.7 + assert guardrails.guardrails_service_url == os.environ["NVIDIA_GUARDRAILS_URL"] + + +def test_init_nemo_guardrails_invalid_temperature(): + from llama_stack.providers.remote.safety.nvidia.nvidia import NeMoGuardrails + + os.environ["NVIDIA_GUARDRAILS_URL"] = "http://nemo.test" + + config = NVIDIASafetyConfig( + guardrails_service_url=os.environ["NVIDIA_GUARDRAILS_URL"], + config_id="test-custom-config-id", + ) + with pytest.raises(ValueError): + NeMoGuardrails(config, "test-model", temperature=0)