erivlis / graphinate

# Plot a graph with custom edge attributes that aren't 'label'

def test_plot_with_custom_edge_attributes(mocker):

    # Arrange

    mock_draw = mocker.patch('graphinate.renderers.matplotlib.draw')

    mock_pyplot = mocker.patch('graphinate.renderers.matplotlib.pyplot')

    mock_ax = mocker.MagicMock()

    mock_fig = mocker.MagicMock()

    mock_pyplot.gca.return_value = mock_ax

    mock_pyplot.gcf.return_value = mock_fig

    graph = nx.Graph(name="Graph with Custom Edge Attributes")

    graph.add_node(1, label="Node 1")

    graph.add_node(2, label="Node 2")

    graph.add_edge(1, 2, label="Edge 1-2", weight=5, type="connection")

    # Act

    from graphinate.renderers.matplotlib import plot

    plot(graph, with_edge_labels=True)

    # Assert

    mock_draw.assert_called_once_with(graph, True, True)

    mock_ax.margins.assert_called_once_with(0.10)

    mock_fig.suptitle.assert_called_once_with("Graph with Custom Edge Attributes")

    mock_fig.tight_layout.assert_called_once()

    mock_pyplot.show.assert_called_once()

# Plot a graph with custom node attributes that aren't 'label'

def test_plot_with_custom_node_attributes(mocker):

    # Arrange

    mock_draw = mocker.patch('graphinate.renderers.matplotlib.draw')

    mock_pyplot = mocker.patch('graphinate.renderers.matplotlib.pyplot')

    mock_ax = mocker.MagicMock()

    mock_fig = mocker.MagicMock()

    mock_pyplot.gca.return_value = mock_ax

    mock_pyplot.gcf.return_value = mock_fig

    graph = nx.Graph(name="Graph with Custom Node Attributes")

    graph.add_node(1, label="Node 1", weight=10, category="A")

    graph.add_node(2, label="Node 2", weight=5, category="B")

    graph.add_edge(1, 2)

    # Act

    from graphinate.renderers.matplotlib import plot

    plot(graph)

    # Assert

    mock_draw.assert_called_once_with(graph, True, False)

    mock_ax.margins.assert_called_once_with(0.10)

    mock_fig.suptitle.assert_called_once_with("Graph with Custom Node Attributes")

    mock_fig.tight_layout.assert_called_once()

    mock_pyplot.show.assert_called_once()

# Plot a graph with custom edge attributes that aren't 'label'