Agent: Use del_key instead of del

monkey/common/utils/code_utils.py

@@ -1,6 +1,6 @@
 import queue
 from collections.abc import MutableSequence
-from typing import Any, List, MutableMapping, TypeVar
+from typing import Any, Dict, List, MutableMapping, Type, TypeVar
 
 T = TypeVar("T")
 
@@ -16,7 +16,7 @@ class abstractstatic(staticmethod):
 
 
 class Singleton(type):
-    _instances = {}
+    _instances: Dict[Type, type] = {}
 
     def __call__(cls, *args, **kwargs):
         if cls not in cls._instances:

monkey/infection_monkey/network/relay/relay_user_handler.py

@@ -5,6 +5,8 @@ from typing import Dict
 
 from egg_timer import EggTimer
 
+from monkey.common.utils.code_utils import del_key
+
 DEFAULT_NEW_CLIENT_TIMEOUT = 3  # Wait up to 3 seconds for potential new clients to connect
 
 
@@ -34,7 +36,7 @@ class RelayUserHandler:
 
         with self._lock:
             if user_address in self._potential_users:
-                del self._potential_users[user_address]
+                del_key(self._potential_users, user_address)
 
             timer = EggTimer()
             self._relay_users[user_address] = RelayUser(user_address, timer)
@@ -59,7 +61,7 @@ class RelayUserHandler:
         """
         with self._lock:
             if user_address in self._relay_users:
-                del self._relay_users[user_address]
+                del_key(self._relay_users, user_address)
 
     def has_potential_users(self) -> bool:
         """