python-graph-visualization/python-wrapper/tests/test_pandas.py at c6989a0495306de50372fefc06e001bbe7588c38 · neo4j/python-graph-visualization · 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
from pandas import DataFrame
from pydantic_extra_types.color import Color

from neo4j_viz.pandas import from_dfs


def test_from_df() -> None:
    nodes = DataFrame(
        {
            "id": [0, 1],
            "caption": ["A", "B"],
            "size": [1337, 42],
            "color": "#FF0000",
        }
    )
    relationships = DataFrame(
        {
            "source": [0, 1],
            "target": [1, 0],
            "caption": ["REL", "REL2"],
        }
    )
    VG = from_dfs(nodes, relationships, node_radius_min_max=(42, 1337))

    assert len(VG.nodes) == 2

    assert VG.nodes[0].id == 0
    assert VG.nodes[0].caption == "A"
    assert VG.nodes[0].size == 1337
    assert VG.nodes[0].color == Color("#ff0000")

    assert VG.nodes[1].id == 1
    assert VG.nodes[1].caption == "B"
    assert VG.nodes[1].size == 42
    assert VG.nodes[0].color == Color("#ff0000")

    assert len(VG.relationships) == 2

    assert VG.relationships[0].source == 0
    assert VG.relationships[0].target == 1
    assert VG.relationships[0].caption == "REL"

    assert VG.relationships[1].source == 1
    assert VG.relationships[1].target == 0
    assert VG.relationships[1].caption == "REL2"


def test_from_dfs() -> None:
    nodes = [
        DataFrame(
            {
                "id": [0],
                "caption": ["A"],
                "size": [1337],
                "color": "#FF0000",
            }
        ),
        DataFrame(
            {
                "id": [1],
                "caption": ["B"],
                "size": [42],
                "color": "#FF0000",
            }
        ),
    ]

    relationships = [
        DataFrame(
            {
                "source": [0],
                "target": [1],
                "caption": ["REL"],
            }
        ),
        DataFrame(
            {
                "source": [1],
                "target": [0],
                "caption": ["REL2"],
            }
        ),
    ]
    VG = from_dfs(nodes, relationships, node_radius_min_max=(42, 1337))

    assert len(VG.nodes) == 2

    assert VG.nodes[0].id == 0
    assert VG.nodes[0].caption == "A"
    assert VG.nodes[0].size == 1337
    assert VG.nodes[0].color == Color("#ff0000")

    assert VG.nodes[1].id == 1
    assert VG.nodes[1].caption == "B"
    assert VG.nodes[1].size == 42
    assert VG.nodes[0].color == Color("#ff0000")

    assert len(VG.relationships) == 2

    assert VG.relationships[0].source == 0
    assert VG.relationships[0].target == 1
    assert VG.relationships[0].caption == "REL"

    assert VG.relationships[1].source == 1
    assert VG.relationships[1].target == 0
    assert VG.relationships[1].caption == "REL2"