Adding remove_role to context and role_handler. Make it possible to a…

…dd and remove roles from the role_context.

mango/agent/role.py

@@ -165,6 +165,15 @@ def add_role(self, role: Role) -> None:
         self._roles.append(role)
         self._role_to_active[role] = True
 
+    def remove_role(self, role: Role) -> None:
+        """Remove a given role
+
+        Args:
+            role ([type]): the role
+        """
+        self._roles.remove(role)
+        del self._role_to_active[role]
+
     @property
     def roles(self) -> List[Role]:
         """Returns all roles
@@ -340,8 +349,21 @@ def add_role(self, role: Role):
 
         :param role: the Role
         """
+        role.bind(self)
         self._role_handler.add_role(role)
 
+        # Setup role
+        role.setup()
+
+    def remove_role(self, role: Role):
+        """Remove a role and call on_stop for clean up
+
+        :param role: the role to remove
+        :type role: Role
+        """
+        self._role_handler.remove_role(role)
+        asyncio.create_task(role.on_stop())
+
     def handle_message(self, content, meta: Dict[str, Any]):
         """Handle an incoming message, delegating it to all applicable subscribers
         for role, message_condition, method, _ in self._message_subs:
@@ -435,20 +457,15 @@ def add_role(self, role: Role):
 
         :param role: the role to add
         """
-        role.bind(self._role_context)
         self._role_context.add_role(role)
 
-        # Setup role
-        role.setup()
-
     def remove_role(self, role: Role):
         """Remove a role permanently from the agent.
 
         :param role: [description]
         :type role: Role
         """
         self._role_context.remove_role(role)
-        asyncio.create_task(role.on_stop())
 
     @property
     def roles(self) -> List[Role]:

tests/unit_tests/role_agent_test.py

@@ -127,6 +127,15 @@ def setup(self):
             self.context.deactivate(r)
 
 
+class TestRole(Role):
+    def __init__(self):
+        self.setup_called = False
+
+    def setup(self):
+        assert self.context is not None
+        self.setup_called = True
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "num_agents,num_containers", [(1, 1), (2, 1), (2, 2), (10, 2), (10, 10)]
@@ -213,3 +222,34 @@ async def test_send_ping_pong_deactivated_pong(num_agents, num_containers):
         await c.shutdown()
 
     assert len(asyncio.all_tasks()) == 1
+
+
+@pytest.mark.asyncio
+async def test_role_add_remove():
+    c = await container_factory.create(addr=("127.0.0.2", 5555))
+    agent = RoleAgent(c)
+    role = TestRole()
+    agent.add_role(role)
+
+    assert agent._role_handler.roles[0] == role
+
+    agent.remove_role(role)
+
+    assert len(agent._role_handler.roles) == 0
+    await c.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_role_add_remove_context():
+    c = await container_factory.create(addr=("127.0.0.2", 5555))
+    agent = RoleAgent(c)
+    role = TestRole()
+    agent._role_context.add_role(role)
+
+    assert role.setup_called
+    assert agent._role_handler.roles[0] == role
+
+    agent._role_context.remove_role(role)
+
+    assert len(agent._role_handler.roles) == 0
+    await c.shutdown()