infragraph/src/tests/test_visualizer.py at 15684607209507537d6b1d27d18bc6eae88e37fb · Keysight/infragraph · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import pytest
import os
import json
import yaml
from infragraph.blueprints.fabrics.clos_fat_tree_fabric import ClosFatTreeFabric
from infragraph.blueprints.devices.generic.server import Server
from infragraph.blueprints.devices.generic.generic_switch import Switch
from infragraph.blueprints.devices.nvidia.dgx import NvidiaDGX
from infragraph.blueprints.devices.nvidia.cx5 import Cx5
from infragraph.blueprints.devices.common.transceiver.qsfp import QSFP
from infragraph.infragraph_service import InfraGraphService
from infragraph.infragraph import *
from infragraph.visualizer.visualize import run_visualizer


def _load_graph_data(output_dir):
    """
    Helper function to assert if the generated graph_data.js has correct number of nodes and edges.
    - Splits the generated json into
        nodes and edges of instances and components
    """
    js_path = os.path.join(output_dir, "js", "graph_data.js")
    with open(js_path, "r") as f:
        content = f.read()
    json_str = content.split("const GRAPH_DATA = ")[1].split(";\n")[0]
    return json.loads(json_str)


@pytest.mark.asyncio
async def test_visualizer_closfabric(tmp_path):
    """
    Test visualizer with a clos fabric infrastructure object directly.
    """
    server = Server()
    switch = Switch(port_count=4)
    clos_fat_tree = ClosFatTreeFabric(switch, server, 3, [])

    output_dir = str(tmp_path / "viz1")
    run_visualizer(infrastructure=clos_fat_tree, hosts="server", switches="switch", output=output_dir)

    assert os.path.exists(os.path.join(output_dir, "index.html")), "HTML file not copied"
    assert os.path.exists(os.path.join(output_dir, "js", "graph_data.js")), "graph_data.js not generated"
    assert os.path.exists(os.path.join(output_dir, "css", "style.css")), "CSS file not copied"

    data = _load_graph_data(output_dir)
    infra_nodes = data["infrastructure.json"]["nodes"]
    infra_edges = data["infrastructure.json"]["edges"]
    assert len(infra_nodes) == 28, "Infrastructure should have 28 total nodes"
    assert len(infra_edges) > 0, "Infrastructure should have edges"


@pytest.mark.asyncio
async def test_composed_devices(tmp_path):
    """
    Test visualizer with composed devices (DGX + CX5) via YAML input.
    """
    cx5 = Cx5(variant="cx5_100g_dual")
    device = NvidiaDGX("dgx_a100", cx5)
    infrastructure = Api().infrastructure()
    infrastructure.devices.append(device).append(cx5)
    infrastructure.instances.add(name=device.name, device=device.name, count=1)

    # dump to yaml
    yaml_path = str(tmp_path / "composed_devices.yaml")
    data = infrastructure.serialize("dict")
    with open(yaml_path, "w") as f:
        yaml.dump(data, f, default_flow_style=False, indent=4)

    output_dir = str(tmp_path / "viz2")
    run_visualizer(input_file=yaml_path, hosts="dgx_a100, cx5_100gbe", switches="", output=output_dir)

    assert os.path.exists(os.path.join(output_dir, "index.html")), "HTML file not copied"
    assert os.path.exists(os.path.join(output_dir, "js", "graph_data.js")), "graph_data.js not generated"
    assert os.path.exists(os.path.join(output_dir, "css", "style.css")), "CSS file not copied"

    data = _load_graph_data(output_dir)
    dgx_nodes = data["dgx_a100.json"]["nodes"]
    dgx_edges = data["dgx_a100.json"]["edges"]
    cx5_nodes = data["cx5_100gbe.json"]["nodes"]
    cx5_edges = data["cx5_100gbe.json"]["edges"]
    assert len(dgx_nodes) == 37, "DGX should have 37 component nodes"
    assert len(dgx_edges) > 0, "DGX should have internal edges"
    assert len(cx5_nodes) == 4, "cx5_100gbe should have 4 nodes"
    assert len(cx5_edges) > 0, "cx5_100gbe should have edges"


@pytest.mark.asyncio
async def test_json_file(tmp_path):
    """
    Test visualizer with DGX + CX5 + QSFP via JSON input.
    """
    qsfp = QSFP("qsfp_dd_400g")
    cx5 = Cx5(variant="cx5_100g_dual", transceiver=qsfp)
    device = NvidiaDGX("dgx_a100", cx5)
    infrastructure = Api().infrastructure()
    infrastructure.devices.append(device).append(cx5).append(qsfp)
    infrastructure.instances.add(name=device.name, device=device.name, count=1)

    service = InfraGraphService()
    service.set_graph(infrastructure)
    service.get_networkx_graph()

    #dump to json
    json_path = str(tmp_path / "infra.json")
    data = infrastructure.serialize("dict")
    with open(json_path, "w") as f:
        json.dump(data, f, indent=2)

    output_dir = str(tmp_path / "viz3")
    run_visualizer(input_file=json_path, hosts="dgx_a100, cx5_100g_dual, qsfp_dd_400g", switches="", output=output_dir)

    assert os.path.exists(os.path.join(output_dir, "index.html")), "HTML file not copied"
    assert os.path.exists(os.path.join(output_dir, "js", "graph_data.js")), "graph_data.js not generated"
    assert os.path.exists(os.path.join(output_dir, "css", "style.css")), "CSS file not copied"

    data = _load_graph_data(output_dir)
    assert len(data) > 1, "Should have infrastructure and at least one device view"


if __name__ == "__main__":
    pytest.main(["-s", __file__])