Fix more tests

Author: JashG

tests/llm_providers/test_langchain_initializer.py

@@ -70,25 +70,27 @@ def test_special_case_called_first(mock_initializers):
 
 def test_chat_completion_called(mock_initializers):
     mock_initializers["special"].return_value = None
+    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = "chat_model"
     result = init_langchain_model("chat-model", "provider", "chat", {})
     assert result == "chat_model"
     mock_initializers["special"].assert_called_once()
+    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_not_called()
-    mock_initializers["text"].assert_not_called()
 
 
 def test_community_chat_called(mock_initializers):
     mock_initializers["special"].return_value = None
+    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = None
     mock_initializers["community"].return_value = "community_model"
     result = init_langchain_model("community-chat", "provider", "chat", {})
     assert result == "community_model"
     mock_initializers["special"].assert_called_once()
+    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_called_once()
-    mock_initializers["text"].assert_not_called()
 
 
 def test_text_completion_called(mock_initializers):
@@ -154,36 +156,39 @@ def test_all_initializers_raise_exceptions(mock_initializers):
 
 def test_duplicate_modes_in_initializer(mock_initializers):
     mock_initializers["special"].return_value = None
+    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = "chat_model"
     result = init_langchain_model("chat-model", "provider", "chat", {})
     assert result == "chat_model"
     mock_initializers["special"].assert_called_once()
+    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_not_called()
-    mock_initializers["text"].assert_not_called()
 
 
 def test_chat_completion_called_when_special_returns_none(mock_initializers):
     mock_initializers["special"].return_value = None
+    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = "chat_model"
     result = init_langchain_model("chat-model", "provider", "chat", {})
     assert result == "chat_model"
     mock_initializers["special"].assert_called_once()
+    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_not_called()
-    mock_initializers["text"].assert_not_called()
 
 
 def test_community_chat_called_when_previous_fail(mock_initializers):
     mock_initializers["special"].return_value = None
+    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = None
     mock_initializers["community"].return_value = "community_model"
     result = init_langchain_model("community-chat", "provider", "chat", {})
     assert result == "community_model"
     mock_initializers["special"].assert_called_once()
+    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_called_once()
-    mock_initializers["text"].assert_not_called()
 
 
 def test_text_completion_called_when_previous_fail(mock_initializers):
@@ -201,12 +206,11 @@ def test_text_completion_called_when_previous_fail(mock_initializers):
 
 def test_text_completion_supports_chat_mode(mock_initializers):
     mock_initializers["special"].return_value = None
-    mock_initializers["chat"].return_value = None
-    mock_initializers["community"].return_value = None
     mock_initializers["text"].return_value = "text_model"
     result = init_langchain_model("text-model", "provider", "chat", {})
     assert result == "text_model"
     mock_initializers["special"].assert_called_once()
-    mock_initializers["chat"].assert_called_once()
-    mock_initializers["community"].assert_called_once()
     mock_initializers["text"].assert_called_once()
+    # Since text returns a value, chat and community are not called
+    mock_initializers["chat"].assert_not_called()
+    mock_initializers["community"].assert_not_called()

tests/test_rails_config.py

@@ -98,7 +98,7 @@ def test_check_prompt_exist_for_self_check_rails():
         ],
     }
     with pytest.raises(
-        ValueError, match="You must provide a `self_check_output` prompt template"
+        ValueError, match="Missing a `self_check_output` prompt template"
     ):
         RailsConfig.check_prompt_exist_for_self_check_rails(values)
 
@@ -353,7 +353,7 @@ def test_validate_rail_prompts_wrong_flow_id_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_input \$model=content_safety` prompt template.",
+            match="Missing a `content_safety_check_input \$model=content_safety` prompt template",
         ):
             _validate_rail_prompts(
                 ["content safety check input $model=content_safety"],
@@ -366,7 +366,7 @@ def test_validate_rail_prompts_wrong_model_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_input \$model=content_safety` prompt template.",
+            match="Missing a `content_safety_check_input \$model=content_safety` prompt template",
         ):
             _validate_rail_prompts(
                 ["content safety check input $model=content_safety"],
@@ -379,7 +379,7 @@ def test_validate_rail_prompts_no_prompt_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_input \$model=content_safety` prompt template.",
+            match="Missing a `content_safety_check_input \$model=content_safety` prompt template",
         ):
             _validate_rail_prompts(
                 ["content safety check input $model=content_safety"],
@@ -395,7 +395,7 @@ def test_content_safety_input_missing_prompt_raises(self):
         """Check Content Safety output rail raises ValueError if we don't have a prompt"""
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_input \$model=content_safety` prompt template.",
+            match="Missing a `content_safety_check_input \$model=content_safety` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -415,7 +415,7 @@ def test_content_safety_output_missing_prompt_raises(self):
         """Check Content Safety output rail raises ValueError if we don't have a prompt"""
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_output \$model=content_safety` prompt template.",
+            match="Missing a `content_safety_check_output \$model=content_safety` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -531,7 +531,7 @@ def test_input_content_safety_no_model_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="No `content_safety` model provided for input flow `content safety check input`",
+            match="Input flow 'content safety check input' references model type 'content_safety' that is not defined",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -556,7 +556,7 @@ def test_input_content_safety_wrong_model_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="No `content_safety` model provided for input flow `content safety check input",
+            match="Input flow 'content safety check input' references model type 'content_safety' that is not defined",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -581,7 +581,7 @@ def test_output_content_safety_no_model_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="No `content_safety` model provided for output flow `content safety check output`",
+            match="Output flow 'content safety check output' references model type 'content_safety' that is not defined",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -606,7 +606,7 @@ def test_output_content_safety_wrong_model_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_output \$model=content_safety` prompt template",
+            match="Missing a `content_safety_check_output \$model=content_safety` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -664,7 +664,7 @@ def test_topic_safety_no_prompt_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="You must provide a `topic_safety_check_input \$model=topic_control` prompt template",
+            match="Missing a `topic_safety_check_input \$model=topic_control` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -688,7 +688,7 @@ def test_topic_safety_no_model_raises(self):
         """Check if we don't provide a topic-safety model we raise a ValueError"""
         with pytest.raises(
             ValueError,
-            match="No `topic_control` model provided for input flow `topic safety check input`",
+            match="Input flow 'topic safety check input' references model type 'topic_control' that is not defined",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -712,7 +712,7 @@ def test_topic_safety_no_model_no_prompt_raises(self):
         """Check a missing model and prompt raises ValueError"""
         with pytest.raises(
             ValueError,
-            match="You must provide a `topic_safety_check_input \$model=topic_control` prompt template",
+            match="Missing a `topic_safety_check_input \$model=topic_control` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -741,7 +741,7 @@ def test_hero_separate_models_no_prompts_raises(self):
 
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_input \$model=my_content_safety` prompt template",
+            match="Missing a `content_safety_check_input \$model=my_content_safety` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -883,7 +883,7 @@ def test_hero_no_prompts_raises(self):
         """Create hero workflow with no prompts. Expect Content Safety input prompt check to fail"""
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_input \$model=content_safety` prompt template",
+            match="Missing a `content_safety_check_input \$model=content_safety` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -923,7 +923,7 @@ def test_hero_no_output_content_safety_prompt_raises(self):
         """Create hero workflow with no prompts. Expect Content Safety input prompt check to fail"""
         with pytest.raises(
             ValueError,
-            match="You must provide a `topic_safety_check_input \$model=your_topic_control` prompt template",
+            match="Missing a `topic_safety_check_input \$model=your_topic_control` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -967,7 +967,7 @@ def test_hero_no_topic_safety_prompt_raises(self):
         """Create hero workflow with no prompts. Expect Content Safety input prompt check to fail"""
         with pytest.raises(
             ValueError,
-            match="You must provide a `topic_safety_check_input \$model=your_topic_control` prompt template",
+            match="Missing a `topic_safety_check_input \$model=your_topic_control` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""
@@ -1013,7 +1013,7 @@ def test_hero_topic_safety_prompt_raises(self):
         """Create hero workflow with no prompts. Expect Content Safety input prompt check to fail"""
         with pytest.raises(
             ValueError,
-            match="You must provide a `content_safety_check_input \$model=content_safety` prompt template",
+            match="Missing a `content_safety_check_input \$model=content_safety` prompt template",
         ):
             _ = RailsConfig.from_content(
                 yaml_content="""