From f4b47f8238310124c555a988e64ad027083a31c4 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Mon, 19 Sep 2022 16:44:37 +0200
Subject: [PATCH] Agent: Use HTTPIslandAPIClient in find_server

---
 monkey/infection_monkey/network/relay/utils.py            | 8 ++++----
 .../infection_monkey/network/relay/test_utils.py          | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py
index 9a60eb7d5..c2ffc64ba 100644
--- a/monkey/infection_monkey/network/relay/utils.py
+++ b/monkey/infection_monkey/network/relay/utils.py
@@ -5,13 +5,13 @@ from ipaddress import IPv4Address
 from typing import Dict, Iterable, Iterator, MutableMapping, Optional
 
 from common.network.network_utils import address_to_ip_port
-from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST
-from infection_monkey.transport import (
-    IslandAPIClient,
+from infection_monkey.island_api_client import (
+    HTTPIslandAPIClient,
     IslandAPIConnectionError,
     IslandAPIError,
     IslandAPITimeoutError,
 )
+from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST
 from infection_monkey.utils.threading import (
     ThreadSafeIterator,
     create_daemon_thread,
@@ -54,7 +54,7 @@ def _check_if_island_server(server: str) -> bool:
     logger.debug(f"Trying to connect to server: {server}")
 
     try:
-        IslandAPIClient(server)
+        HTTPIslandAPIClient(server)
 
         return True
     except IslandAPIConnectionError as err:
diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py
index ad22e0a4c..c4e69bd1a 100644
--- a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py
+++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py
@@ -1,8 +1,8 @@
 import pytest
 import requests_mock
 
+from infection_monkey.island_api_client import IslandAPIConnectionError
 from infection_monkey.network.relay.utils import find_server
-from infection_monkey.transport import IslandAPIConnectionError
 
 SERVER_1 = "1.1.1.1:12312"
 SERVER_2 = "2.2.2.2:4321"