From 46573c9f2dba5dfdb8185dc6f27207d3a40ee63b Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Thu, 1 Sep 2022 19:24:13 +0000 Subject: [PATCH] UT: Add tests for RelayUserHandler --- .../network/relay/test_relay_user_handler.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py new file mode 100644 index 000000000..34aaf6bc7 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py @@ -0,0 +1,23 @@ +from ipaddress import IPv4Address + +from monkey.infection_monkey.network.relay import RelayUserHandler + + +def test_potential_users_added(): + user_address = IPv4Address("0.0.0.0") + handler = RelayUserHandler() + + assert len(handler.get_potential_users()) == 0 + handler.add_potential_user(user_address) + assert len(handler.get_potential_users()) == 1 + assert user_address in handler.get_potential_users() + + +def test_potential_user_removed_on_matching_user_added(): + user_address = IPv4Address("0.0.0.0") + handler = RelayUserHandler() + + handler.add_potential_user(user_address) + handler.add_relay_user(user_address) + + assert len(handler.get_potential_users()) == 0