fixed tests

Author: rkritika1508

backend/app/tests/test_toxicity_hub_validators.py

@@ -37,7 +37,7 @@ def test_build_with_default_policies(self):
     def test_build_with_explicit_policies(self):
         config = LlamaGuard7BSafetyValidatorConfig(
             type="llamaguard_7b",
-            policies=["O1", "O2"],
+            policies=["no_violence_hate", "no_sexual_content"],
         )
 
         with patch(_LLAMAGUARD_PATCH) as mock_validator:
@@ -56,7 +56,14 @@ def test_build_with_empty_policies_list(self):
         assert kwargs["policies"] == []
 
     def test_build_with_all_policy_codes(self):
-        all_policies = ["O1", "O2", "O3", "O4", "O5", "O6"]
+        all_policies = [
+            "no_violence_hate",
+            "no_sexual_content",
+            "no_criminal_planning",
+            "no_guns_and_illegal_weapons",
+            "no_illegal_drugs",
+            "no_encourage_self_harm",
+        ]
         config = LlamaGuard7BSafetyValidatorConfig(
             type="llamaguard_7b", policies=all_policies
         )
@@ -65,11 +72,11 @@ def test_build_with_all_policy_codes(self):
             config.build()
 
         _, kwargs = mock_validator.call_args
-        assert kwargs["policies"] == all_policies
+        assert kwargs["policies"] == ["O1", "O2", "O3", "O4", "O5", "O6"]
 
     def test_build_with_single_policy(self):
         config = LlamaGuard7BSafetyValidatorConfig(
-            type="llamaguard_7b", policies=["O3"]
+            type="llamaguard_7b", policies=["no_criminal_planning"]
         )
 
         with patch(_LLAMAGUARD_PATCH) as mock_validator:
@@ -78,6 +85,15 @@ def test_build_with_single_policy(self):
         _, kwargs = mock_validator.call_args
         assert kwargs["policies"] == ["O3"]
 
+    def test_build_with_invalid_policy_raises(self):
+        config = LlamaGuard7BSafetyValidatorConfig(
+            type="llamaguard_7b", policies=["O1"]
+        )
+
+        with patch(_LLAMAGUARD_PATCH):
+            with pytest.raises(ValueError, match="Unknown policy"):
+                config.build()
+
     def test_build_returns_validator_instance(self):
         config = LlamaGuard7BSafetyValidatorConfig(type="llamaguard_7b")