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 <melbeher@redhat.com>

tests/unit/providers/nvidia/test_safety.py

@@ -5,53 +5,61 @@
 # 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):
+class TestNVIDIASafetyAdapter(NVIDIASafetyAdapter):
-    def setUp(self):
+    """Test implementation that provides the required shield_store."""
+
+    def __init__(self, config: NVIDIASafetyConfig, shield_store):
+        super().__init__(config)
+        self.shield_store = shield_store
+
+
+@pytest.fixture
+def nvidia_adapter():
+    """Set up the NVIDIASafetyAdapter for testing."""
     os.environ["NVIDIA_GUARDRAILS_URL"] = "http://nemo.test"
 
     # Initialize the adapter
-        self.config = NVIDIASafetyConfig(
+    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
 
-        # Mock the HTTP request methods
+    # Create a mock shield store that implements the ShieldStore protocol
-        self.guardrails_post_patcher = patch(
+    shield_store = AsyncMock()
-            "llama_stack.providers.remote.safety.nvidia.nvidia.NeMoGuardrails._guardrails_post"
+    shield_store.get_shield = AsyncMock()
-        )
-        self.mock_guardrails_post = self.guardrails_post_patcher.start()
-        self.mock_guardrails_post.return_value = {"status": "allowed"}
 
-    def tearDown(self):
+    adapter = TestNVIDIASafetyAdapter(config=config, shield_store=shield_store)
-        """Clean up after each test."""
-        self.guardrails_post_patcher.stop()
 
-    @pytest.fixture(autouse=True)
+    return adapter
-    def inject_fixtures(self, run_async):
-        self.run_async = run_async
 
-    def _assert_request(
+
-        self,
+@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 _assert_request(
     mock_call: MagicMock,
     expected_url: str,
     expected_headers: dict[str, str] | None = None,
     expected_json: dict[str, Any] | None = None,
-    ) -> None:
+) -> None:
     """
     Helper method to verify request details in mock API calls.
 
@@ -80,42 +88,51 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
             else:
                 assert call_args[1]["json"][key] == value
 
-    def test_register_shield_with_valid_id(self):
+
+async def test_register_shield_with_valid_id(nvidia_adapter):
+    adapter = nvidia_adapter
+
     shield = Shield(
         provider_id="nvidia",
-            type="shield",
+        type=ResourceType.shield,
         identifier="test-shield",
         provider_resource_id="test-model",
     )
 
     # Register the shield
-        self.run_async(self.adapter.register_shield(shield))
+    await adapter.register_shield(shield)
+
+
+async def test_register_shield_without_id(nvidia_adapter):
+    adapter = nvidia_adapter
 
-    def test_register_shield_without_id(self):
     shield = Shield(
         provider_id="nvidia",
-            type="shield",
+        type=ResourceType.shield,
         identifier="test-shield",
         provider_resource_id="",
     )
 
     # Register the shield should raise a ValueError
-        with self.assertRaises(ValueError):
+    with pytest.raises(ValueError):
-            self.run_async(self.adapter.register_shield(shield))
+        await adapter.register_shield(shield)
+
+
+async def test_run_shield_allowed(nvidia_adapter, mock_guardrails_post):
+    adapter = nvidia_adapter
 
-    def test_run_shield_allowed(self):
     # Set up the shield
     shield_id = "test-shield"
     shield = Shield(
         provider_id="nvidia",
-            type="shield",
+        type=ResourceType.shield,
         identifier=shield_id,
         provider_resource_id="test-model",
     )
-        self.shield_store.get_shield.return_value = shield
+    adapter.shield_store.get_shield.return_value = shield
 
     # Mock Guardrails API response
-        self.mock_guardrails_post.return_value = {"status": "allowed"}
+    mock_guardrails_post.return_value = {"status": "allowed"}
 
     # Run the shield
     messages = [
@@ -123,17 +140,17 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
         CompletionMessage(
             role="assistant",
             content="I'm doing well, thank you for asking!",
-                stop_reason="end_of_message",
+            stop_reason=StopReason.end_of_message,
             tool_calls=[],
         ),
     ]
-        result = self.run_async(self.adapter.run_shield(shield_id, messages))
+    result = await adapter.run_shield(shield_id, messages)
 
     # Verify the shield store was called
-        self.shield_store.get_shield.assert_called_once_with(shield_id)
+    adapter.shield_store.get_shield.assert_called_once_with(shield_id)
 
     # Verify the Guardrails API was called correctly
-        self.mock_guardrails_post.assert_called_once_with(
+    mock_guardrails_post.assert_called_once_with(
         path="/v1/guardrail/checks",
         data={
             "model": shield_id,
@@ -157,19 +174,22 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
     assert isinstance(result, RunShieldResponse)
     assert result.violation is None
 
-    def test_run_shield_blocked(self):
+
+async def test_run_shield_blocked(nvidia_adapter, mock_guardrails_post):
+    adapter = nvidia_adapter
+
     # Set up the shield
     shield_id = "test-shield"
     shield = Shield(
         provider_id="nvidia",
-            type="shield",
+        type=ResourceType.shield,
         identifier=shield_id,
         provider_resource_id="test-model",
     )
-        self.shield_store.get_shield.return_value = shield
+    adapter.shield_store.get_shield.return_value = shield
 
     # Mock Guardrails API response
-        self.mock_guardrails_post.return_value = {"status": "blocked", "rails_status": {"reason": "harmful_content"}}
+    mock_guardrails_post.return_value = {"status": "blocked", "rails_status": {"reason": "harmful_content"}}
 
     # Run the shield
     messages = [
@@ -177,17 +197,17 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
         CompletionMessage(
             role="assistant",
             content="I'm doing well, thank you for asking!",
-                stop_reason="end_of_message",
+            stop_reason=StopReason.end_of_message,
             tool_calls=[],
         ),
     ]
-        result = self.run_async(self.adapter.run_shield(shield_id, messages))
+    result = await adapter.run_shield(shield_id, messages)
 
     # Verify the shield store was called
-        self.shield_store.get_shield.assert_called_once_with(shield_id)
+    adapter.shield_store.get_shield.assert_called_once_with(shield_id)
 
     # Verify the Guardrails API was called correctly
-        self.mock_guardrails_post.assert_called_once_with(
+    mock_guardrails_post.assert_called_once_with(
         path="/v1/guardrail/checks",
         data={
             "model": shield_id,
@@ -214,37 +234,43 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
     assert result.violation.violation_level == ViolationLevel.ERROR
     assert result.violation.metadata == {"reason": "harmful_content"}
 
-    def test_run_shield_not_found(self):
+
+async def test_run_shield_not_found(nvidia_adapter, mock_guardrails_post):
+    adapter = nvidia_adapter
+
     # Set up shield store to return None
     shield_id = "non-existent-shield"
-        self.shield_store.get_shield.return_value = None
+    adapter.shield_store.get_shield.return_value = None
 
     messages = [
         UserMessage(role="user", content="Hello, how are you?"),
     ]
 
-        with self.assertRaises(ValueError):
+    with pytest.raises(ValueError):
-            self.run_async(self.adapter.run_shield(shield_id, messages))
+        await adapter.run_shield(shield_id, messages)
 
     # Verify the shield store was called
-        self.shield_store.get_shield.assert_called_once_with(shield_id)
+    adapter.shield_store.get_shield.assert_called_once_with(shield_id)
 
     # Verify the Guardrails API was not called
-        self.mock_guardrails_post.assert_not_called()
+    mock_guardrails_post.assert_not_called()
+
+
+async def test_run_shield_http_error(nvidia_adapter, mock_guardrails_post):
+    adapter = nvidia_adapter
 
-    def test_run_shield_http_error(self):
     shield_id = "test-shield"
     shield = Shield(
         provider_id="nvidia",
-            type="shield",
+        type=ResourceType.shield,
         identifier=shield_id,
         provider_resource_id="test-model",
     )
-        self.shield_store.get_shield.return_value = shield
+    adapter.shield_store.get_shield.return_value = shield
 
     # Mock Guardrails API to raise an exception
     error_msg = "API Error: 500 Internal Server Error"
-        self.mock_guardrails_post.side_effect = Exception(error_msg)
+    mock_guardrails_post.side_effect = Exception(error_msg)
 
     # Running the shield should raise an exception
     messages = [
@@ -252,18 +278,18 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
         CompletionMessage(
             role="assistant",
             content="I'm doing well, thank you for asking!",
-                stop_reason="end_of_message",
+            stop_reason=StopReason.end_of_message,
             tool_calls=[],
         ),
     ]
-        with self.assertRaises(Exception) as context:
+    with pytest.raises(Exception) as exc_info:
-            self.run_async(self.adapter.run_shield(shield_id, messages))
+        await adapter.run_shield(shield_id, messages)
 
     # Verify the shield store was called
-        self.shield_store.get_shield.assert_called_once_with(shield_id)
+    adapter.shield_store.get_shield.assert_called_once_with(shield_id)
 
     # Verify the Guardrails API was called correctly
-        self.mock_guardrails_post.assert_called_once_with(
+    mock_guardrails_post.assert_called_once_with(
         path="/v1/guardrail/checks",
         data={
             "model": shield_id,
@@ -283,11 +309,14 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
         },
     )
     # Verify the exception message
-        assert error_msg in str(context.exception)
+    assert error_msg in str(exc_info.value)
 
-    def test_init_nemo_guardrails(self):
+
+def test_init_nemo_guardrails():
     from llama_stack.providers.remote.safety.nvidia.nvidia import NeMoGuardrails
 
+    os.environ["NVIDIA_GUARDRAILS_URL"] = "http://nemo.test"
+
     test_config_id = "test-custom-config-id"
     config = NVIDIASafetyConfig(
         guardrails_service_url=os.environ["NVIDIA_GUARDRAILS_URL"],
@@ -314,12 +343,15 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
     assert guardrails.temperature == 0.7
     assert guardrails.guardrails_service_url == os.environ["NVIDIA_GUARDRAILS_URL"]
 
-    def test_init_nemo_guardrails_invalid_temperature(self):
+
+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 self.assertRaises(ValueError):
+    with pytest.raises(ValueError):
         NeMoGuardrails(config, "test-model", temperature=0)