class TestNode(TestCase): def setUp(self): self.node = Node('x') self.agent1 = Agent(1) self.agent2 = Agent(2) self.y = Node('y') self.colorY = 'red' self.z = Node('z') self.colorZ = 'green' def test_init(self): self.assertEqual('x', self.node.id) self.assertSetEqual(set(), self.node.agents) self.assertDictEqual({}, self.node.neighbors) def test_addAgentsGetAgents(self): self.node.addAgents([self.agent1]) self.assertSetEqual({self.agent1}, self.node.getAgents()) def test_removeAgents(self): self.node.addAgents([self.agent1, self.agent2]) self.node.removeAgents([self.agent1]) self.assertSetEqual({self.agent2}, self.node.getAgents()) def test_addEdgesGetEdgeColorsGetNeighbor(self): self.node.addEdges({self.colorY: self.y, self.colorZ: self.z}) self.assertItemsEqual([self.colorY, self.colorZ], self.node.getEdgeColors()) self.assertEqual(self.y, self.node.getNeighbor(self.colorY)) self.assertEqual(self.z, self.node.getNeighbor(self.colorZ)) def test_removeEdges(self): self.node.addEdges({self.colorY: self.y, self.colorZ: self.z}) self.node.removeEdges([self.colorY]) self.assertItemsEqual([self.colorZ], self.node.getEdgeColors()) self.assertEqual(self.z, self.node.getNeighbor(self.colorZ))