import pytest from unittest.mock import patch, AsyncMock, MagicMock from src.graph import build_tree_structure, tree_to_dot, create_hierarchical_view from src.network import make_payload, get_graph from src.viz import create_pyvis_network from app import get_id_from_input import asyncio # sample graph for testing SAMPLE_GRAPH = { "nodes": { 1: {"name": "Root", "year": 2000, "advisors": [], "institution": "A"}, 2: {"name": "Child", "year": 2020, "advisors": [1], "institution": "B"}, 3: {"name": "Grandchild", "year": 2040, "advisors": [2], "institution": "C"}, 4: {"name": "Unrelated", "year": 1990, "advisors": [], "institution": "D"}, } } def test_build_tree_structure_basic(): tree = build_tree_structure(SAMPLE_GRAPH, 1) # only the root node is included when root_id=1 assert list(tree.keys()) == [1] assert tree[1]["depth"] == 0 assert tree[1]["advisors"] == [] assert tree[1]["children"] == [] def test_tree_to_dot_output(): dot = tree_to_dot(SAMPLE_GRAPH) assert "digraph G" in dot assert '"1" [label="Root (2000)"' in dot assert '"2" [label="Child (2020)"' in dot assert '"1" -> "2";' in dot def test_create_hierarchical_view_nodes_edges(): nodes, edges = create_hierarchical_view(SAMPLE_GRAPH, 1) node_ids = {n.id for n in nodes} # only the root node is present assert node_ids == {"1"} assert edges == [] def test_make_payload_structure(): payload = make_payload(123) assert payload["kind"] == "build-graph" assert payload["options"]["reportingCallback"] is True assert payload["startNodes"][0]["recordId"] == 123 @pytest.mark.asyncio async def test_get_graph_mocks_websocket(): payload = make_payload(1) fake_graph = {"kind": "graph", "payload": {"nodes": {1: {"name": "Root", "advisors": []}}}} # patch websockets.connect to yield a mock websocket with patch("src.network.websockets.connect", new_callable=AsyncMock) as mock_connect: mock_ws = AsyncMock() mock_connect.return_value.__aenter__.return_value = mock_ws # simulate .recv() returning a graph message mock_ws.recv = AsyncMock(side_effect=[ # first message: progress '{"kind": "progress", "payload": {"queued": 1, "fetching": 0, "done": 0}}', # second message: graph '{"kind": "graph", "payload": {"nodes": {"1": {"name": "Root", "advisors": []}}}}' ]) result = await get_graph(payload) assert "nodes" in result assert 1 in result["nodes"] or "1" in result["nodes"] def test_create_pyvis_network_html(): html = create_pyvis_network(SAMPLE_GRAPH, 1) assert "