test: Update coverage improvement test for tests/test_generate_answer_node.py

Author: codebeaver-ai[bot]

tests/test_generate_answer_node.py

@@ -268,3 +268,87 @@ def test_init_chat_ollama_format():
         "dummy_input", ["output"], node_config=node_config_with_schema
     )
     assert node2.llm_model.format == "dummy_schema_json"
+
+import time
+import json
+from requests.exceptions import Timeout
+
+# New dummy chain classes and dummy llms for testing invoke_with_timeout and ChatBedrock behavior.
+class DummySlowChain:
+    """A dummy chain that simulates a slow response by returning quickly but using patched time."""
+    def invoke(self, inputs):
+        return "response"
+
+class DummyChainException:
+    """A dummy chain that simulates an exception during invoke."""
+    def invoke(self, inputs):
+        raise Exception("dummy error")
+
+class DummyChatBedrock:
+    """A dummy ChatBedrock chain that returns a predefined answer. It should be used when output_parser is None."""
+    def __call__(self, *args, **kwargs):
+        return {"content": "bedrock answer"}
+
+    def __or__(self, other):
+        return self
+
+# New tests to increase test coverage.
+def test_invoke_with_timeout_exceeded(dummy_node):
+    """
+    Test that invoke_with_timeout raises a Timeout exception when the chain.invoke 
+    takes longer than the specified timeout.
+    """
+    # Set dummy_node.timeout to a known value (from dummy_node, timeout=1 is used)
+    original_time = time.time
+    # Create a list of times such that the first call returns start_time, 
+    # and the next call returns start_time + timeout + 1.
+    call_times = [100, 100 + dummy_node.timeout + 1]
+    def fake_time():
+        return call_times.pop(0) if call_times else original_time()
+    time_time_backup = time.time
+    time.time = fake_time
+    try:
+        with pytest.raises(Timeout):
+            dummy_node.invoke_with_timeout(DummySlowChain(), {"key": "value"}, dummy_node.timeout)
+    finally:
+        time.time = time_time_backup
+
+def test_invoke_with_timeout_exception_propagates(dummy_node):
+    """
+    Test that invoke_with_timeout properly propagates exceptions raised by the chain.invoke call.
+    """
+    with pytest.raises(Exception) as excinfo:
+        dummy_node.invoke_with_timeout(DummyChainException(), {"key": "value"}, dummy_node.timeout)
+    assert "dummy error" in str(excinfo.value)
+
+@pytest.mark.parametrize("content_key", ["relevant_chunks", "parsed_doc", "doc", "content"])
+def test_process_with_alternative_content_keys(dummy_node_for_process, content_key):
+    """
+    Test that process() successfully uses alternative content keys.
+    """
+    state = {"user_prompt": "What is the answer?", content_key: "Alternative content"}
+    dummy_node_for_process.chain = DummyChain()
+    dummy_node_for_process.invoke_with_timeout = lambda chain, inputs, timeout: chain.invoke(inputs)
+    result_state = dummy_node_for_process.process(state)
+    assert result_state["output"] == {"content": "successful answer"}
+
+def test_execute_single_chunk_with_bedrock():
+    """
+    Test execute() for the single chunk branch using a DummyChatBedrock instance.
+    This verifies that if the llm_model is a ChatBedrock, no output_parser is applied.
+    """
+    node_config = {"llm_model": DummyChatBedrock(), "verbose": False, "timeout": 480}
+    node = GenerateAnswerNode("dummy_input & doc", ["output"], node_config=node_config)
+    node.logger = DummyLogger()
+    node.get_input_keys = lambda state: ["dummy_input", "doc"]
+    state = {"dummy_input": "What is the answer?", "doc": ["Only one chunk bedrock"]}
+
+    # simulate a simple chain invocation (in the single-chunk branch)
+    def fake_invoke_with_timeout(chain, inputs, timeout):
+        if "question" in inputs:
+            return {"content": "bedrock answer"}
+        return {"content": "unexpected response"}
+
+    node.invoke_with_timeout = fake_invoke_with_timeout
+    output_state = node.execute(state)
+    assert output_state["output"] == {"content": "bedrock answer"}