From c601f2214a015632a15d77d3f653cf3922299906 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 6 Sep 2022 16:39:41 +0000 Subject: [PATCH] Agent: Use del_key instead of del --- monkey/common/utils/code_utils.py | 4 ++-- monkey/infection_monkey/network/relay/relay_user_handler.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/monkey/common/utils/code_utils.py b/monkey/common/utils/code_utils.py index 52fdcbaac..4ddff97e2 100644 --- a/monkey/common/utils/code_utils.py +++ b/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: diff --git a/monkey/infection_monkey/network/relay/relay_user_handler.py b/monkey/infection_monkey/network/relay/relay_user_handler.py index e73049a63..6af2c841e 100644 --- a/monkey/infection_monkey/network/relay/relay_user_handler.py +++ b/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: """