tests/test_functionality.py

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 "<html" in html.lower()
    assert "Root" in html
    assert "Child" in html

@pytest.mark.parametrize("val,expected", [
    ("123", 123),
    ("0", 0),
    ("notanint", None),
    ("", None),
    (" 42 ", 42),
])
def test_get_id_from_input(val, expected):
    assert get_id_from_input(val) == expected