def test_to_dict_when_entities_are_linked_in_loop(self): graph = Graph() e1 = Entity(1, 'E1') e2 = Entity(2, 'E2') graph.add_entity(e1) graph.add_entity(e2) graph.link_entities(e1, e2) graph.link_entities(e2, e1) self.assertDictEqual(graph.entities, {1: e1, 2: e2}) json_dict = graph.to_dict() self.assert_graph_dict( json_dict, { 'entities': [{ 'entity_id': 1, 'name': 'E1' }, { 'entity_id': 2, 'name': 'E2' }], 'links': [{ 'from': 1, 'to': 2 }, { 'from': 2, 'to': 1 }], })
def test_to_dict_when_entity_has_multiple_links(self): graph = Graph() e1 = Entity(1, 'E1') e2 = Entity(2, 'E2') e3 = Entity(3, 'E3') graph.add_entity(e1) graph.add_entity(e2) graph.add_entity(e3) graph.link_entities(e1, e3) graph.link_entities(e2, e3) self.assertDictEqual(graph.entities, {1: e1, 2: e2, 3: e3}) json_dict = graph.to_dict() self.assert_graph_dict( json_dict, { 'entities': [{ 'entity_id': 1, 'name': 'E1' }, { 'entity_id': 2, 'name': 'E2' }, { 'entity_id': 3, 'name': 'E3' }], 'links': [{ 'from': 1, 'to': 3 }, { 'from': 2, 'to': 3 }], })
class TestGraphLinkEntities(TestCase, AssertEntityMixin): def setUp(self): self.graph = Graph() self.entity_1 = Entity(1, 'E1') self.entity_2 = Entity(2, 'E2') self.graph.add_entity(self.entity_1) self.graph.add_entity(self.entity_2) def test_link_entities_when_arguments_are_none(self): self.graph.link_entities(self.entity_1, None) self.graph.link_entities(None, self.entity_2) self.graph.link_entities(None, None) self.assert_links(self.entity_1) self.assert_links(self.entity_2) def test_link_entities(self): self.graph.link_entities(self.entity_1, self.entity_2) self.assert_links(self.entity_1, successors=set([self.entity_2])) self.assert_links(self.entity_2, predecessors=set([self.entity_1])) def test_link_entities_by_id(self): self.graph.link_entities_by_id(1, 2) self.assert_links(self.entity_1, successors=set([self.entity_2])) self.assert_links(self.entity_2, predecessors=set([self.entity_1]))