tests/test_langgraph_nodes.py

import pytest
from unittest.mock import patch, AsyncMock, MagicMock
from app import synthesize_responses # Removed generate_visualization from here as it's the SUT for one test
from insight_state import InsightFlowState
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_openai import ChatOpenAI # For spec
from openai import AsyncOpenAI # For spec of the client if we were to patch app.openai_async_client
# Import app to allow patching its global openai_async_client
import app as application_module 

# Define a plausible system prompt for the synthesizer for use in assertion
SYNTHESIZER_SYSTEM_PROMPT_TEMPLATE = """You are a master synthesizer AI. Your task is to integrate the following diverse perspectives into a single, coherent, and insightful response. Ensure that the final synthesis is well-structured, easy to understand, and accurately reflects the nuances of each provided viewpoint. Do not simply list the perspectives; weave them together.

Perspectives:
{formatted_perspectives}

Synthesized Response:"""

@pytest.mark.asyncio
async def test_synthesize_responses_node():
    """
    Test the synthesize_responses node to ensure it calls the llm_synthesizer
    with the correct prompt and updates the state.
    """
    initial_state = InsightFlowState(
        query="Test Query",
        selected_personas=["analytical", "scientific"],
        persona_responses={
            "analytical": "Analytical perspective text.",
            "scientific": "Scientific perspective text."
        },
        synthesized_response=None,
        current_step_name="execute_persona_tasks" # Previous step
    )

    mock_synthesized_text = "This is the beautifully synthesized response from the LLM."
    
    # Expected formatted perspectives string
    expected_perspectives_text = ("- Perspective from analytical: Analytical perspective text.\n"
                                  "- Perspective from scientific: Scientific perspective text.")
    expected_final_prompt_content = SYNTHESIZER_SYSTEM_PROMPT_TEMPLATE.format(formatted_perspectives=expected_perspectives_text)

    # Create a mock for the entire llm_synthesizer object in app.py
    mock_llm_synthesizer_replacement = MagicMock(spec=ChatOpenAI)
    # Mock its ainvoke method
    mock_llm_synthesizer_replacement.ainvoke = AsyncMock(return_value=AIMessage(content=mock_synthesized_text))

    # Patch app.llm_synthesizer to be our mock_llm_synthesizer_replacement
    with patch('app.llm_synthesizer', new=mock_llm_synthesizer_replacement) as mock_synthesizer_in_app:
        # Call the actual node function
        # synthesize_responses (from app) will now use the patched llm_synthesizer (mock_synthesizer_in_app)
        output_state = await synthesize_responses(initial_state.copy())

    # Assertions
    mock_synthesizer_in_app.ainvoke.assert_called_once()
    args, kwargs = mock_synthesizer_in_app.ainvoke.call_args
    called_messages = args[0] # The 'messages' argument
    
    assert len(called_messages) == 1 # Expecting a single system message for simplicity here, or system + human
    # For now, let's assume the prompt is constructed as a single SystemMessage containing everything
    # This might evolve when we implement the actual prompt construction in app.py
    assert isinstance(called_messages[0], SystemMessage) 
    assert called_messages[0].content == expected_final_prompt_content
    
    assert output_state["synthesized_response"] == mock_synthesized_text
    assert output_state["current_step_name"] == "generate_visualization"
    assert output_state["persona_responses"] == initial_state["persona_responses"] # Should not change 

@pytest.mark.asyncio
@patch('app.generate_dalle_image') # Corrected: Patch where it's used by app.generate_visualization
@patch('app.openai_async_client') 
@patch('app.generate_mermaid_code') # <--- ADDED PATCH for mermaid generation utility
async def test_generate_visualization_node_creates_dalle_and_mermaid(
    mock_app_generate_mermaid_code, # <--- ADDED mock argument
    mock_app_openai_client, 
    mock_app_level_generate_dalle_image_func # Renamed to reflect it's app's view of the function
):
    """
    Test the generate_visualization node to ensure it calls generate_dalle_image
    and generate_mermaid_code, and updates the state with image URL and mermaid code.
    """
    initial_synthesized_response = "This is a detailed synthesized response about complex topics."
    initial_state = InsightFlowState(
        query="Test Query for Visuals",
        synthesized_response=initial_synthesized_response,
        persona_responses={}, # Explicitly initialize for completeness
        visualization_image_url=None,
        visualization_code=None,
        current_step_name="synthesize_responses" 
    )

    expected_dalle_url = "https://dalle.example.com/generated_image.png"
    mock_app_level_generate_dalle_image_func.return_value = expected_dalle_url
    
    expected_mermaid_output = "graph LR; A-->B; C-->D;" # Mocked mermaid output
    mock_app_generate_mermaid_code.return_value = expected_mermaid_output
    
    # Ensure llm_mermaid_generator is available in the application_module for the test context
    # If app.llm_mermaid_generator is None during the test, the call might be skipped.
    # We assume it's initialized by initialize_configurations(), which should run before/during app import or be called by on_chat_start.
    # For node tests, if direct LLM calls are made, ensure they are properly mocked or the LLMs are available.
    # Here, generate_mermaid_code (the util) takes the LLM as an arg, which app.generate_visualization supplies.
    # So, we need to make sure app.llm_mermaid_generator is a mock or a real object for the call.
    # Let's patch app.llm_mermaid_generator to be a MagicMock for this test to ensure it's passed correctly.
    with patch.object(application_module, 'llm_mermaid_generator', new_callable=MagicMock) as mock_llm_mermaid_in_app:
        output_state = await application_module.generate_visualization(initial_state.copy())

    # DALL-E Assertions
    mock_app_level_generate_dalle_image_func.assert_called_once()
    called_dalle_kwargs = mock_app_level_generate_dalle_image_func.call_args.kwargs
    assert called_dalle_kwargs.get('prompt') == f"A hand-drawn style visual note or sketch representing the key concepts of: {initial_synthesized_response}"
    assert called_dalle_kwargs.get('client') is mock_app_openai_client
    assert output_state["visualization_image_url"] == expected_dalle_url

    # Mermaid Assertions
    mock_app_generate_mermaid_code.assert_called_once_with(
        initial_synthesized_response, 
        mock_llm_mermaid_in_app # Assert it was called with the (mocked) app.llm_mermaid_generator
    )
    assert output_state["visualization_code"] == expected_mermaid_output 
    
    # General State Assertions
    assert output_state["current_step_name"] == "present_results"
    assert output_state["synthesized_response"] == initial_synthesized_response
    assert output_state["persona_responses"] == initial_state["persona_responses"] # Should not change 

@pytest.mark.asyncio
@patch('app.cl.user_session.get')
@patch('app.cl.Image') # Mock for cl.Image class
@patch('app.cl.Text')  # <--- ADDED PATCH for cl.Text class
@patch('app.cl.Message') # Mock for cl.Message class
async def test_present_results_node_sends_all_content(
    mock_message_class, 
    mock_text_class,    # <--- ADDED mock argument
    mock_image_class, 
    mock_user_session_get,
):
    """Test present_results node sends synthesized response, visuals, and perspectives."""
    
    # Configure app.cl.user_session.get mock
    def user_session_get_side_effect(key, default=None):
        if key == "show_visualization":
            return True
        if key == "show_perspectives":
            return True
        # Fallback for other keys used by the node if any (e.g. in future)
        # Currently, present_results from app.py doesn't use default values from cl.user_session.get for these booleans directly in its logic
        # It implies they should exist. The main app's on_chat_start sets them.
        # For this test, we only care about show_visualization and show_perspectives.
        return default
    mock_user_session_get.side_effect = user_session_get_side_effect

    # Configure cl.Message mock (the class itself)
    # Each time cl.Message() is called, it returns a new mock_message_instance.
    # This instance needs an async 'send' method.
    # To allow multiple calls to cl.Message and check them, we need Message instances to be distinct if their send methods are checked individually.
    # A simpler way for now is to check the call_args_list of mock_message_class.
    # The send method is on the instance, so we make the class return an instance that has an async send.
    async def mock_send(*args, **kwargs):
        pass # Mock async send
    
    # Store all created message instances to check their send calls individually if needed
    created_message_mocks = [] 
    def message_constructor_side_effect(*args, **kwargs):
        instance = MagicMock(name=f"MockClMessageInstance_{len(created_message_mocks)}")
        instance.send = AsyncMock(wraps=mock_send) # each instance gets its own send mock
        instance.content = kwargs.get("content")
        instance.elements = kwargs.get("elements")
        created_message_mocks.append(instance)
        return instance

    mock_message_class.side_effect = message_constructor_side_effect

    # Configure cl.Image mock (the class itself)
    mock_image_instance_returned_by_constructor = MagicMock(name="MockClImageInstance")
    mock_image_class.return_value = mock_image_instance_returned_by_constructor

    # Configure cl.Text mock (the class itself)
    mock_text_instance_returned_by_constructor = MagicMock(name="MockClTextInstance")
    mock_text_class.return_value = mock_text_instance_returned_by_constructor

    initial_state = application_module.InsightFlowState(
        query="Test Query for Presentation",
        synthesized_response="Final synthesized answer.",
        persona_responses={"analytical": "Analytical point.", "creative": "Creative idea."},
        visualization_image_url="http://example.com/dalle.png",
        visualization_code="graph TD; X-->Y;",
        current_step_name="generate_visualization" # Previous step
    )

    output_state = await application_module.present_results(initial_state.copy())

    # Assertions for user_session.get calls
    mock_user_session_get.assert_any_call("show_visualization")
    mock_user_session_get.assert_any_call("show_perspectives")

    # Assertions for cl.Message calls
    # Expected calls: 1 for synthesized, 1 for DALL-E, 1 for Mermaid, 2 for perspectives = 5 calls
    assert mock_message_class.call_count == 5
    # Ensure each created message mock had its send method called once
    for msg_mock_instance in created_message_mocks:
        msg_mock_instance.send.assert_called_once()

    all_message_constructor_calls = mock_message_class.call_args_list

    # 1. Synthesized response
    call_synthesized_kwargs = all_message_constructor_calls[0].kwargs
    assert call_synthesized_kwargs.get('content') == "Final synthesized answer."
    assert not call_synthesized_kwargs.get('elements') 

    # 2. DALL-E Image
    call_dalle_kwargs = all_message_constructor_calls[1].kwargs
    mock_image_class.assert_called_once_with(
        url="http://example.com/dalle.png", 
        name="dalle_visualization", 
        display="inline",
        size="large" # Default in app.py
    )
    assert call_dalle_kwargs.get('elements') == [mock_image_instance_returned_by_constructor]
    # Content for image message could be None or specific text. Let's assume it's empty if not specified.
    assert call_dalle_kwargs.get('content') == "" # Or check for a specific title if app.py adds one

    # 3. Mermaid Diagram
    call_mermaid_kwargs = all_message_constructor_calls[2].kwargs
    mock_text_class.assert_called_once_with(
        content="graph TD; X-->Y;", # The raw mermaid code from initial_state
        mime_type="text/mermaid",
        name="generated_diagram",
        display="inline"
    )
    assert call_mermaid_kwargs.get('elements') == [mock_text_instance_returned_by_constructor]
    assert call_mermaid_kwargs.get('content') == "" # Expecting empty content for message with element

    # 4. Persona Perspectives (order within perspectives might not be guaranteed due to dict iteration)
    perspective_contents_sent = set()
    perspective_contents_sent.add(all_message_constructor_calls[3].kwargs.get('content'))
    perspective_contents_sent.add(all_message_constructor_calls[4].kwargs.get('content'))

    expected_perspective1_content = "**Perspective from analytical:**\nAnalytical point."
    expected_perspective2_content = "**Perspective from creative:**\nCreative idea."
    assert expected_perspective1_content in perspective_contents_sent
    assert expected_perspective2_content in perspective_contents_sent
    
    # Check state update
    assert output_state["current_step_name"] == "results_presented"