python-graph-visualization/python-wrapper/tests/test_pandas.py at f9731fc5e45c4ba296982764909a37ab03cd8c2b · 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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
from pandas import DataFrame
from pydantic_extra_types.color import Color

from neo4j_viz.node import Node
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", "instrument": ["piano", "guitar"]}
    )
    relationships = DataFrame(
        {
            "source": [0, 1],
            "target": [1, 0],
            "caption": ["REL", "REL2"],
            "weight": [1.0, 2.0],
        }
    )
    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[0].properties == {"instrument": "piano"}

    assert VG.nodes[1].id == 1
    assert VG.nodes[1].caption == "B"
    assert VG.nodes[1].size == 42
    assert VG.nodes[1].color == Color("#ff0000")
    assert VG.nodes[1].properties == {"instrument": "guitar"}

    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[0].properties == {"weight": 1.0}

    assert VG.relationships[1].source == 1
    assert VG.relationships[1].target == 0
    assert VG.relationships[1].caption == "REL2"
    assert VG.relationships[1].properties == {"weight": 2.0}


def test_from_rel_dfs() -> None:
    relationships = [
        DataFrame(
            {
                "source": [0, 1],
                "target": [1, 0],
                "caption": ["REL", "REL2"],
                "weight": [1.0, 2.0],
            }
        ),
        DataFrame(
            {
                "source": [2, 3],
                "target": [1, 0],
                "caption": ["REL", "REL2"],
                "weight": [1.0, 2.0],
            }
        ),
    ]
    VG = from_dfs(None, relationships)

    assert len(VG.relationships) == 4
    assert VG.nodes == [Node(id=id) for id in [0, 1, 2, 3]]


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"