# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 from unittest.mock import MagicMock, patch import pytest import requests from haystack.core.errors import PipelineDrawingError from haystack.core.pipeline import Pipeline from haystack.core.pipeline.draw import _to_mermaid_image, _to_mermaid_text from haystack.testing.sample_components import AddFixedValue, Double @pytest.mark.flaky(reruns=5, reruns_delay=5) @pytest.mark.integration def test_to_mermaid_image(): pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") image_data = _to_mermaid_image(pipe.graph) # We just verify we received some data as testing the actual image is not reliable assert image_data @patch("haystack.core.pipeline.draw.requests") def test_to_mermaid_image_does_not_edit_graph(mock_requests): pipe = Pipeline() pipe.add_component("comp1", AddFixedValue(add=3)) pipe.add_component("comp2", Double()) pipe.connect("comp1.result", "comp2.value") pipe.connect("comp2.value", "comp1.value") mock_requests.get.return_value = MagicMock(status_code=200) expected_pipe = pipe.to_dict() _to_mermaid_image(pipe.graph) assert expected_pipe == pipe.to_dict() def test_to_mermaid_image_failing_request(tmp_path): pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") pipe.connect("comp2", "comp1") with patch("haystack.core.pipeline.draw.requests.get") as mock_get: def raise_for_status(self): raise requests.HTTPError() mock_response = MagicMock() mock_response.status_code = 429 mock_response.content = '{"error": "too many requests"}' mock_response.raise_for_status = raise_for_status mock_get.return_value = mock_response with pytest.raises(PipelineDrawingError, match="There was an issue with https://mermaid.ink"): _to_mermaid_image(pipe.graph) def test_to_mermaid_text(): pipe = Pipeline() pipe.add_component("comp1", AddFixedValue(add=3)) pipe.add_component("comp2", Double()) pipe.connect("comp1.result", "comp2.value") pipe.connect("comp2.value", "comp1.value") init_params = {"theme": "neutral"} text = _to_mermaid_text(pipe.graph, init_params) assert ( text == """ %%{ init: {'theme': 'neutral'} }%% graph TD; comp1["comp1
AddFixedValue

Optional inputs:
  • add (Optional[int])
"]:::component -- "result -> value
int" --> comp2["comp2
Double"]:::component comp2["comp2
Double"]:::component -- "value -> value
int" --> comp1["comp1
AddFixedValue

Optional inputs:
  • add (Optional[int])
"]:::component classDef component text-align:center; """ ) def test_to_mermaid_text_does_not_edit_graph(): pipe = Pipeline() pipe.add_component("comp1", AddFixedValue(add=3)) pipe.add_component("comp2", Double()) pipe.connect("comp1.result", "comp2.value") pipe.connect("comp2.value", "comp1.value") expected_pipe = pipe.to_dict() init_params = {"theme": "neutral"} _to_mermaid_text(pipe.graph, init_params) assert expected_pipe == pipe.to_dict() @pytest.mark.integration @pytest.mark.parametrize( "params", [ {"format": "img", "type": "png", "theme": "dark"}, {"format": "svg", "theme": "forest"}, {"format": "pdf", "fit": True, "theme": "neutral"}, ], ) def test_to_mermaid_image_valid_formats(params): # Test valid formats pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") image_data = _to_mermaid_image(pipe.graph, params=params) assert image_data # Ensure some data is returned def test_to_mermaid_image_invalid_format(): # Test invalid format pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") with pytest.raises(ValueError, match="Invalid image format:"): _to_mermaid_image(pipe.graph, params={"format": "invalid_format"}) @pytest.mark.integration def test_to_mermaid_image_missing_theme(): # Test default theme (neutral) pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") params = {"format": "img"} image_data = _to_mermaid_image(pipe.graph, params=params) assert image_data # Ensure some data is returned def test_to_mermaid_image_invalid_scale(): # Test invalid scale pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") with pytest.raises(ValueError, match="Scale must be a number between 1 and 3."): _to_mermaid_image(pipe.graph, params={"format": "img", "scale": 5}) def test_to_mermaid_image_scale_without_dimensions(): # Test scale without width/height pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") with pytest.raises(ValueError, match="Scale is only allowed when width or height is set."): _to_mermaid_image(pipe.graph, params={"format": "img", "scale": 2}) @patch("haystack.core.pipeline.draw.requests.get") def test_to_mermaid_image_server_error(mock_get): # Test server failure pipe = Pipeline() pipe.add_component("comp1", Double()) pipe.add_component("comp2", Double()) pipe.connect("comp1", "comp2") def raise_for_status(self): raise requests.HTTPError() mock_response = MagicMock() mock_response.status_code = 500 mock_response.content = '{"error": "server error"}' mock_response.raise_for_status = raise_for_status mock_get.return_value = mock_response with pytest.raises(PipelineDrawingError, match="There was an issue with https://mermaid.ink"): _to_mermaid_image(pipe.graph) def test_to_mermaid_image_invalid_server_url(): # Test invalid server URL pipe = Pipeline() pipe.add_component("comp1", AddFixedValue(add=3)) pipe.add_component("comp2", Double()) pipe.connect("comp1.result", "comp2.value") pipe.connect("comp2.value", "comp1.value") server_url = "https://invalid.server" with pytest.raises(PipelineDrawingError, match=f"There was an issue with {server_url}"): _to_mermaid_image(pipe.graph, server_url=server_url)