Files
2026-07-13 12:36:30 +08:00

66 lines
2.2 KiB
Python

import unittest
import easygraph as eg
class TestEgoGraph(unittest.TestCase):
def setUp(self):
self.simple_graph = eg.Graph()
self.simple_graph.add_edges_from([(0, 1), (1, 2), (2, 3), (3, 4)])
self.directed_graph = eg.DiGraph()
self.directed_graph.add_edges_from([(0, 1), (1, 2), (2, 3)])
self.weighted_graph = eg.Graph()
self.weighted_graph.add_edges_from(
[(0, 1, {"weight": 1}), (1, 2, {"weight": 2}), (2, 3, {"weight": 3})]
)
self.disconnected_graph = eg.Graph()
self.disconnected_graph.add_edges_from([(0, 1), (2, 3)])
self.single_node_graph = eg.Graph()
self.single_node_graph.add_node(42)
def test_simple_graph_radius_1(self):
ego = eg.functions.community.ego_graph(self.simple_graph, 2, radius=1)
self.assertSetEqual(set(ego.nodes), {1, 2, 3})
def test_simple_graph_radius_2(self):
ego = eg.functions.community.ego_graph(self.simple_graph, 2, radius=2)
self.assertSetEqual(set(ego.nodes), {0, 1, 2, 3, 4})
def test_directed_graph(self):
ego = eg.functions.community.ego_graph(self.directed_graph, 1, radius=1)
self.assertSetEqual(set(ego.nodes), {1, 2})
def test_weighted_graph_with_distance(self):
ego = eg.functions.community.ego_graph(
self.weighted_graph, 0, radius=2, distance="weight"
)
self.assertSetEqual(set(ego.nodes), {0, 1})
def test_disconnected_graph(self):
ego = eg.functions.community.ego_graph(self.disconnected_graph, 0, radius=1)
self.assertSetEqual(set(ego.nodes), {0, 1})
def test_single_node_graph(self):
ego = eg.functions.community.ego_graph(self.single_node_graph, 42, radius=1)
self.assertSetEqual(set(ego.nodes), {42})
def test_center_false(self):
ego = eg.functions.community.ego_graph(
self.simple_graph, 2, radius=1, center=False
)
self.assertSetEqual(set(ego.nodes), {1, 3})
def test_empty_graph(self):
G = eg.Graph()
G.add_node("x")
ego = eg.functions.community.ego_graph(G, "x", radius=1)
self.assertSetEqual(set(ego.nodes), {"x"})
if __name__ == "__main__":
unittest.main()