add tests

fix

Author: Pouyanpi

tests/test_content_safety_actions.py

@@ -13,18 +13,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
-# conftest.py
 import pytest
 
 from nemoguardrails.library.content_safety.actions import (
+    DEFAULT_REFUSAL_MESSAGES,
+    SUPPORTED_LANGUAGES,
+    _detect_language,
+    _get_refusal_message,
     content_safety_check_input,
     content_safety_check_output,
     content_safety_check_output_mapping,
+    detect_language,
 )
 from tests.utils import FakeLLM
 
+try:
+    import fast_langdetect  # noqa
+
+    HAS_FAST_LANGDETECT = True
+except ImportError:
+    HAS_FAST_LANGDETECT = False
+
+requires_fast_langdetect = pytest.mark.skipif(not HAS_FAST_LANGDETECT, reason="fast-langdetect not installed")
+
 
 @pytest.fixture
 def fake_llm():
@@ -150,3 +163,148 @@ def test_content_safety_check_output_mapping_default():
     """Test content_safety_check_output_mapping defaults to allowed=False when key is missing."""
     result = {"policy_violations": []}
     assert content_safety_check_output_mapping(result) is False
+
+
+@requires_fast_langdetect
+class TestDetectLanguage:
+    @pytest.mark.parametrize(
+        "text,expected_lang",
+        [
+            ("Hello, how are you today?", "en"),
+            ("Hola, ¿cómo estás hoy?", "es"),
+            ("你好，你今天好吗？", "zh"),
+            ("Guten Tag, wie geht es Ihnen?", "de"),
+            ("Bonjour, comment allez-vous?", "fr"),
+            ("こんにちは、お元気ですか？", "ja"),
+        ],
+        ids=["english", "spanish", "chinese", "german", "french", "japanese"],
+    )
+    def test_detect_language(self, text, expected_lang):
+        assert _detect_language(text) == expected_lang
+
+    def test_detect_language_empty_string(self):
+        result = _detect_language("")
+        assert result is None or result == "en"
+
+    def test_detect_language_import_error(self):
+        with patch.dict("sys.modules", {"fast_langdetect": None}):
+            import nemoguardrails.library.content_safety.actions as actions_module
+
+            _original_detect_language = actions_module._detect_language
+
+            def patched_detect_language(text):
+                try:
+                    raise ImportError("No module named 'fast_langdetect'")
+                except ImportError:
+                    return None
+
+            with patch.object(actions_module, "_detect_language", patched_detect_language):
+                result = actions_module._detect_language("Hello")
+                assert result is None
+
+    def test_detect_language_exception(self):
+        with patch("fast_langdetect.detect", side_effect=Exception("Detection failed")):
+            result = _detect_language("Hello")
+            assert result is None
+
+
+class TestGetRefusalMessage:
+    @pytest.mark.parametrize("lang", list(SUPPORTED_LANGUAGES))
+    def test_default_messages(self, lang):
+        result = _get_refusal_message(lang, None)
+        assert result == DEFAULT_REFUSAL_MESSAGES[lang]
+
+    def test_custom_message_used_when_available(self):
+        custom = {"en": "Custom refusal", "es": "Rechazo personalizado"}
+        assert _get_refusal_message("en", custom) == "Custom refusal"
+        assert _get_refusal_message("es", custom) == "Rechazo personalizado"
+
+    def test_unsupported_lang_falls_back_to_english(self):
+        assert _get_refusal_message("xyz", None) == DEFAULT_REFUSAL_MESSAGES["en"]
+        assert _get_refusal_message("xyz", {"en": "Custom fallback"}) == "Custom fallback"
+
+    def test_lang_not_in_custom_uses_default(self):
+        custom = {"en": "Custom English"}
+        assert _get_refusal_message("es", custom) == DEFAULT_REFUSAL_MESSAGES["es"]
+
+
+@requires_fast_langdetect
+class TestDetectLanguageAction:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "user_message,expected_lang",
+        [
+            ("Hello, how are you?", "en"),
+            ("Hola, ¿cómo estás?", "es"),
+            ("你好", "zh"),
+        ],
+        ids=["english", "spanish", "chinese"],
+    )
+    async def test_detect_language_action(self, user_message, expected_lang):
+        context = {"user_message": user_message}
+        result = await detect_language(context=context, config=None)
+        assert result["language"] == expected_lang
+        assert result["refusal_message"] == DEFAULT_REFUSAL_MESSAGES[expected_lang]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "context",
+        [None, {"user_message": ""}],
+        ids=["no_context", "empty_message"],
+    )
+    async def test_detect_language_action_defaults_to_english(self, context):
+        result = await detect_language(context=context, config=None)
+        assert result["language"] == "en"
+        assert result["refusal_message"] == DEFAULT_REFUSAL_MESSAGES["en"]
+
+    @pytest.mark.asyncio
+    async def test_detect_language_action_unsupported_language_falls_back_to_english(self):
+        with patch(
+            "nemoguardrails.library.content_safety.actions._detect_language",
+            return_value="xyz",
+        ):
+            context = {"user_message": "some text"}
+            result = await detect_language(context=context, config=None)
+            assert result["language"] == "en"
+            assert result["refusal_message"] == DEFAULT_REFUSAL_MESSAGES["en"]
+
+    @pytest.mark.asyncio
+    async def test_detect_language_action_with_config_custom_messages(self):
+        mock_config = MagicMock()
+        mock_config.rails.config.content_safety.multilingual.refusal_messages = {
+            "en": "Custom: Cannot help",
+            "es": "Personalizado: No puedo ayudar",
+        }
+
+        context = {"user_message": "Hello"}
+        result = await detect_language(context=context, config=mock_config)
+        assert result["language"] == "en"
+        assert result["refusal_message"] == "Custom: Cannot help"
+
+    @pytest.mark.asyncio
+    async def test_detect_language_action_with_config_no_multilingual(self):
+        mock_config = MagicMock()
+        mock_config.rails.config.content_safety.multilingual = None
+
+        context = {"user_message": "Hello"}
+        result = await detect_language(context=context, config=mock_config)
+        assert result["language"] == "en"
+        assert result["refusal_message"] == DEFAULT_REFUSAL_MESSAGES["en"]
+
+
+class TestSupportedLanguagesAndDefaults:
+    def test_supported_languages_count(self):
+        assert len(SUPPORTED_LANGUAGES) == 9
+
+    def test_supported_languages_contents(self):
+        expected = {"en", "es", "zh", "de", "fr", "hi", "ja", "ar", "th"}
+        assert SUPPORTED_LANGUAGES == expected
+
+    def test_default_refusal_messages_has_all_supported_languages(self):
+        for lang in SUPPORTED_LANGUAGES:
+            assert lang in DEFAULT_REFUSAL_MESSAGES
+
+    def test_default_refusal_messages_are_non_empty(self):
+        for _lang, message in DEFAULT_REFUSAL_MESSAGES.items():
+            assert message
+            assert len(message) > 0

tests/test_rails_config.py

@@ -23,7 +23,9 @@
 
 from nemoguardrails.llm.prompts import TaskPrompt
 from nemoguardrails.rails.llm.config import (
+    ContentSafetyConfig,
     Model,
+    MultilingualConfig,
     RailsConfig,
     _get_flow_model,
     _validate_rail_prompts,
@@ -1015,3 +1017,76 @@ def test_hero_topic_safety_prompt_raises(self):
                     content: Verify the user input is on-topic
                 """
             )
+
+
+class TestMultilingualConfig:
+    def test_defaults(self):
+        config = MultilingualConfig()
+        assert config.enabled is False
+        assert config.refusal_messages is None
+
+    def test_with_custom_messages(self):
+        custom = {"en": "Custom", "es": "Personalizado"}
+        config = MultilingualConfig(enabled=True, refusal_messages=custom)
+        assert config.enabled is True
+        assert config.refusal_messages == custom
+
+
+class TestContentSafetyConfigModel:
+    def test_defaults(self):
+        config = ContentSafetyConfig()
+        assert config.multilingual.enabled is False
+        assert config.multilingual.refusal_messages is None
+
+    def test_with_multilingual(self):
+        custom = {"en": "Custom"}
+        config = ContentSafetyConfig(multilingual=MultilingualConfig(enabled=True, refusal_messages=custom))
+        assert config.multilingual.enabled is True
+        assert config.multilingual.refusal_messages == custom
+
+
+class TestMultilingualConfigInRailsConfig:
+    BASE_YAML = """
+        models:
+          - type: content_safety
+            engine: nim
+            model: nvidia/llama-3.1-nemoguard-8b-content-safety
+        rails:
+          {rails_config}
+          input:
+            flows:
+              - content safety check input $model=content_safety
+        prompts:
+          - task: content_safety_check_input $model=content_safety
+            content: Check content safety
+    """
+
+    def test_multilingual_disabled_by_default(self):
+        config = RailsConfig.from_content(yaml_content=self.BASE_YAML.format(rails_config=""))
+        assert config.rails.config.content_safety.multilingual.enabled is False
+
+    def test_multilingual_enabled_with_custom_messages(self):
+        rails_config = """
+          config:
+            content_safety:
+              multilingual:
+                enabled: true
+                refusal_messages:
+                  en: "Custom English"
+                  es: "Personalizado"
+        """
+        config = RailsConfig.from_content(yaml_content=self.BASE_YAML.format(rails_config=rails_config))
+        assert config.rails.config.content_safety.multilingual.enabled is True
+        assert config.rails.config.content_safety.multilingual.refusal_messages["en"] == "Custom English"
+        assert config.rails.config.content_safety.multilingual.refusal_messages["es"] == "Personalizado"
+
+    def test_multilingual_enabled_no_custom_messages(self):
+        rails_config = """
+          config:
+            content_safety:
+              multilingual:
+                enabled: true
+        """
+        config = RailsConfig.from_content(yaml_content=self.BASE_YAML.format(rails_config=rails_config))
+        assert config.rails.config.content_safety.multilingual.enabled is True
+        assert config.rails.config.content_safety.multilingual.refusal_messages is None