1import networkx as nx
2
3
4def small_ego_G():
5    """The sample network from https://arxiv.org/pdf/1310.6753v1.pdf"""
6    edges = [
7        ("a", "b"),
8        ("a", "c"),
9        ("b", "c"),
10        ("b", "d"),
11        ("b", "e"),
12        ("b", "f"),
13        ("c", "d"),
14        ("c", "f"),
15        ("c", "h"),
16        ("d", "f"),
17        ("e", "f"),
18        ("f", "h"),
19        ("h", "j"),
20        ("h", "k"),
21        ("i", "j"),
22        ("i", "k"),
23        ("j", "k"),
24        ("u", "a"),
25        ("u", "b"),
26        ("u", "c"),
27        ("u", "d"),
28        ("u", "e"),
29        ("u", "f"),
30        ("u", "g"),
31        ("u", "h"),
32        ("u", "i"),
33        ("u", "j"),
34        ("u", "k"),
35    ]
36    G = nx.Graph()
37    G.add_edges_from(edges)
38
39    return G
40
41
42class TestDispersion:
43    def test_article(self):
44        """our algorithm matches article's"""
45        G = small_ego_G()
46        disp_uh = nx.dispersion(G, "u", "h", normalized=False)
47        disp_ub = nx.dispersion(G, "u", "b", normalized=False)
48        assert disp_uh == 4
49        assert disp_ub == 1
50
51    def test_results_length(self):
52        """there is a result for every node"""
53        G = small_ego_G()
54        disp = nx.dispersion(G)
55        disp_Gu = nx.dispersion(G, "u")
56        disp_uv = nx.dispersion(G, "u", "h")
57        assert len(disp) == len(G)
58        assert len(disp_Gu) == len(G) - 1
59        assert type(disp_uv) is float
60
61    def test_impossible_things(self):
62        G = nx.karate_club_graph()
63        disp = nx.dispersion(G)
64        for u in disp:
65            for v in disp[u]:
66                assert disp[u][v] >= 0
67