Repository: Azure/WALinuxAgent
Pull request #1656 (2019-10-06 21:58:31)

Make obtaining wireserver endpoint consistent
## Description
There are current several consumers (direct and indirect) of the Wireserver ip address. Some have the default value hardcoded. Some obtain it when the is_dhcp_available method returns False. In one case, a consumer expects get_wireserver_endpoint method to return None in a scenario to fallback to the default, known ip address.

This PR removes the default ip address from the return of is_dhcp_available and has get_wireserver_endpoint return it when reading from the file fails, which is the case when the file does not exist.

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] Except for special cases involving multiple contributors, the PR is started from a fork of the main repository, not a branch.
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made and Travis.CI is passing.

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/azure/walinuxagent/1656)
<!-- Reviewable:end -->


codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656?src=pr&el=h1) Report
> Merging [#1656](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656?src=pr&el=desc) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/6a970c90d3258a8fcfc24a5303633bb2182af5bd?src=pr&el=desc) will **increase** coverage by `0.09%`.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656/graphs/tree.svg?width=650&token=ae90E2KYCw&height=150&src=pr)](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           develop    #1656      +/-   ##
===========================================
+ Coverage    66.92%   67.02%   +0.09%     
===========================================
  Files           78       78              
  Lines        11263    11263              
  Branches      1575     1575              
===========================================
+ Hits          7538     7549      +11     
+ Misses        3402     3393       -9     
+ Partials       323      321       -2
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/protocol/util.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC91dGlsLnB5) | `78.53% <0%> (+5.75%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656?src=pr&el=footer). Last update [6a970c9...71909ad](https://codecov.io/gh/Azure/WALinuxAgent/pull/1656?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

jacwil: Relates to #233
narrieta: LGTM; will merge after current release.

Patch:
diff --git a/azurelinuxagent/common/dhcp.py b/azurelinuxagent/common/dhcp.py
index fc2cb649..59749656 100644
--- a/azurelinuxagent/common/dhcp.py
+++ b/azurelinuxagent/common/dhcp.py
@@ -19,8 +19,6 @@ import socket
 import array
 import time
 import azurelinuxagent.common.logger as logger
-import azurelinuxagent.common.utils.shellutil as shellutil
-from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.textutil import hex_dump, hex_dump2, \
     hex_dump3, \
     compare_bytes, str_to_ord, \
@@ -29,9 +27,10 @@ from azurelinuxagent.common.utils.textutil import hex_dump, hex_dump2, \
 from azurelinuxagent.common.exception import DhcpError
 from azurelinuxagent.common.osutil import get_osutil
 
+
 # the kernel routing table representation of 168.63.129.16
 KNOWN_WIRESERVER_IP_ENTRY = '10813FA8'
-KNOWN_WIRESERVER_IP = '168.63.129.16'
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP
 
 
 def get_dhcp_handler():
@@ -152,10 +151,10 @@ class DhcpHandler(object):
         """
         Check if DHCP is available
         """
-        (dhcp_available, endpoint) =  self.osutil.is_dhcp_available()
+        dhcp_available =  self.osutil.is_dhcp_available()
         if not dhcp_available:
             logger.info("send_dhcp_req: DHCP not available")
-            self.endpoint = endpoint
+            self.endpoint = KNOWN_WIRESERVER_IP
             return
 
         """
diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py
index 7ccc5344..b99bca43 100644
--- a/azurelinuxagent/common/osutil/default.py
+++ b/azurelinuxagent/common/osutil/default.py
@@ -1072,7 +1072,7 @@ class DefaultOSUtil(object):
                       chk_err=False)
 
     def is_dhcp_available(self):
-        return (True, '')
+        return True
 
     def is_dhcp_enabled(self):
         return False
diff --git a/azurelinuxagent/common/osutil/iosxe.py b/azurelinuxagent/common/osutil/iosxe.py
index 3152ada8..0d45b0e2 100644
--- a/azurelinuxagent/common/osutil/iosxe.py
+++ b/azurelinuxagent/common/osutil/iosxe.py
@@ -63,7 +63,7 @@ class IosxeOSUtil(DefaultOSUtil):
         DefaultOSUtil.openssl_to_openssh(self, input_file, output_file)
 
     def is_dhcp_available(self):
-        return (False, '168.63.129.16')
+        return False
 
     def get_instance_id(self):
         '''
diff --git a/azurelinuxagent/common/protocol/util.py b/azurelinuxagent/common/protocol/util.py
index 692b5489..dab5a26a 100644
--- a/azurelinuxagent/common/protocol/util.py
+++ b/azurelinuxagent/common/protocol/util.py
@@ -35,9 +35,9 @@ from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.dhcp import get_dhcp_handler
 from azurelinuxagent.common.protocol.ovfenv import OvfEnv
 from azurelinuxagent.common.protocol.wire import WireProtocol
-from azurelinuxagent.common.protocol.metadata import MetadataProtocol, \
-                                                     METADATA_ENDPOINT
-from azurelinuxagent.common.utils.restutil import IOErrorCounter
+from azurelinuxagent.common.protocol.metadata import MetadataProtocol
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP, \
+                                                  IOErrorCounter
 
 OVF_FILE_NAME = "ovf-env.xml"
 TAG_FILE_NAME = "useMetadataEndpoint.tag"
@@ -64,14 +64,16 @@ def get_protocol_util():
 
 
 class ProtocolUtil(object):
-
     """
     ProtocolUtil handles initialization for protocol instance. 2 protocol types
     are invoked, wire protocol and metadata protocols.
     """
+
     def __init__(self):
         self.lock = threading.Lock()
+        self.lock_wireserver_endpoint = threading.Lock()
         self.protocol = None
+        self.endpoint = None
         self.osutil = get_osutil()
         self.dhcp_handler = get_dhcp_handler()
 
@@ -95,7 +97,7 @@ class ProtocolUtil(object):
         try:
             ovfxml = fileutil.read_file(ovf_file_path_on_dvd, remove_bom=True)
             ovfenv = OvfEnv(ovfxml)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise ProtocolError("[CopyOvfEnv] Error reading file "
                                 "{0}: {1}".format(ovf_file_path_on_dvd,
                                                   ustr(e)))
@@ -105,7 +107,7 @@ class ProtocolUtil(object):
                             PASSWORD_REPLACEMENT,
                             ovfxml)
             fileutil.write_file(ovf_file_path, ovfxml)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise ProtocolError("[CopyOvfEnv] Error writing file "
                                 "{0}: {1}".format(ovf_file_path,
                                                   ustr(e)))
@@ -114,7 +116,7 @@ class ProtocolUtil(object):
             if os.path.isfile(tag_file_path_on_dvd):
                 logger.info("Found {0} in provisioning ISO", TAG_FILE_NAME)
                 shutil.copyfile(tag_file_path_on_dvd, tag_file_path)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise ProtocolError("[CopyOvfEnv] Error copying file "
                                 "{0} to {1}: {2}".format(tag_file_path,
                                                          tag_file_path,
@@ -147,24 +149,77 @@ class ProtocolUtil(object):
             conf.get_lib_dir(),
             PROTOCOL_FILE_NAME)
 
+    def _get_wireserver_endpoint_file_path(self):
+        return os.path.join(
+            conf.get_lib_dir(),
+            ENDPOINT_FILE_NAME)
+
     def _get_tag_file_path(self):
         return os.path.join(
             conf.get_lib_dir(),
             TAG_FILE_NAME)
 
     def get_wireserver_endpoint(self):
+        self.lock_wireserver_endpoint.acquire()
+
         try:
-            file_path = os.path.join(conf.get_lib_dir(), ENDPOINT_FILE_NAME)
-            return fileutil.read_file(file_path)
-        except IOError as e:
-            raise OSUtilError(ustr(e))
+            if self.endpoint:
+                return self.endpoint
+
+            file_path = self._get_wireserver_endpoint_file_path()
+            if os.path.isfile(file_path):
+                try:
+                    self.endpoint = fileutil.read_file(file_path)
+
+                    if self.endpoint:
+                        logger.info("WireServer endpoint {0} read from file", self.endpoint)
+                        return self.endpoint
+
+                    logger.error("[GetWireserverEndpoint] Unexpected empty file {0}: {1}", file_path, str(e))
+                except (IOError, OSError) as e:
+                    logger.error("[GetWireserverEndpoint] Error reading file {0}: {1}", file_path, str(e))
+            else:
+                logger.error("[GetWireserverEndpoint] Missing file {0}", file_path)
+
+            self.endpoint = KNOWN_WIRESERVER_IP
+            logger.info("Using hardcoded Wireserver endpoint {0}", self.endpoint)
+
+            return self.endpoint
+        finally:
+            self.lock_wireserver_endpoint.release()
 
     def _set_wireserver_endpoint(self, endpoint):
+        self.lock_wireserver_endpoint.acquire()
         try:
-            file_path = os.path.join(conf.get_lib_dir(), ENDPOINT_FILE_NAME)
+            self.endpoint = endpoint
+            file_path = self._get_wireserver_endpoint_file_path()
             fileutil.write_file(file_path, endpoint)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise OSUtilError(ustr(e))
+        finally:
+            self.lock_wireserver_endpoint.release()
+
+    def _clear_wireserver_endpoint(self):
+        """
+        Cleanup previous saved wireserver endpoint.
+        """
+        self.lock_wireserver_endpoint.acquire()
+
+        try:
+            self.endpoint = None
+            endpoint_file_path = self._get_wireserver_endpoint_file_path()
+            if not os.path.isfile(endpoint_file_path):
+                return
+
+            try:
+                os.remove(endpoint_file_path)
+            except (IOError, OSError) as e:
+                # Ignore file-not-found errors (since the file is being removed)
+                if e.errno == errno.ENOENT:
+                    return
+                logger.error("Failed to clear wiresever endpoint: {0}", e)
+        finally:
+            self.lock_wireserver_endpoint.release()
 
     def _detect_wire_protocol(self):
         endpoint = self.dhcp_handler.endpoint
@@ -172,7 +227,7 @@ class ProtocolUtil(object):
             '''
             Check if DHCP can be used to get the wire protocol endpoint
             '''
-            (dhcp_available, conf_endpoint) =  self.osutil.is_dhcp_available()
+            dhcp_available = self.osutil.is_dhcp_available()
             if dhcp_available:
                 logger.info("WireServer endpoint is not found. Rerun dhcp handler")
                 try:
@@ -183,11 +238,6 @@ class ProtocolUtil(object):
             else:
                 logger.info("_detect_wire_protocol: DHCP not available")
                 endpoint = self.get_wireserver_endpoint()
-                if endpoint == None:
-                    endpoint = conf_endpoint
-                    logger.info("Using hardcoded WireServer endpoint {0}", endpoint)
-                else:
-                    logger.info("WireServer endpoint {0} read from file", endpoint)
 
         try:
             protocol = WireProtocol(endpoint)
@@ -195,7 +245,7 @@ class ProtocolUtil(object):
             self._set_wireserver_endpoint(endpoint)
             return protocol
         except ProtocolError as e:
-            logger.info("WireServer is not responding. Reset endpoint")
+            logger.info("WireServer is not responding. Reset dhcp endpoint")
             self.dhcp_handler.endpoint = None
             self.dhcp_handler.skip_cache = True
             raise e
@@ -254,14 +304,15 @@ class ProtocolUtil(object):
         protocol_file_path = self._get_protocol_file_path()
         try:
             fileutil.write_file(protocol_file_path, protocol_name)
-        except IOError as e:
+        except (IOError, OSError) as e:
             logger.error("Failed to save protocol endpoint: {0}", e)
 
     def clear_protocol(self):
         """
-        Cleanup previous saved endpoint.
+        Cleanup previous saved protocol endpoint.
         """
-        logger.info("Clean protocol")
+        logger.info("Clean protocol and wireserver endpoint")
+        self._clear_wireserver_endpoint()
         self.protocol = None
         protocol_file_path = self._get_protocol_file_path()
         if not os.path.isfile(protocol_file_path):
@@ -269,7 +320,7 @@ class ProtocolUtil(object):
 
         try:
             os.remove(protocol_file_path)
-        except IOError as e:
+        except (IOError, OSError) as e:
             # Ignore file-not-found errors (since the file is being removed)
             if e.errno == errno.ENOENT:
                 return
diff --git a/azurelinuxagent/common/utils/restutil.py b/azurelinuxagent/common/utils/restutil.py
index addba28b..fd78b969 100644
--- a/azurelinuxagent/common/utils/restutil.py
+++ b/azurelinuxagent/common/utils/restutil.py
@@ -98,13 +98,13 @@ HTTP_USER_AGENT_HEALTH = "{0}+health".format(HTTP_USER_AGENT)
 INVALID_CONTAINER_CONFIGURATION = "InvalidContainerConfiguration"
 REQUEST_ROLE_CONFIG_FILE_NOT_FOUND = "RequestRoleConfigFileNotFound"
 
-DEFAULT_PROTOCOL_ENDPOINT = '168.63.129.16'
+KNOWN_WIRESERVER_IP = '168.63.129.16'
 HOST_PLUGIN_PORT = 32526
 
 
 class IOErrorCounter(object):
     _lock = threading.RLock()
-    _protocol_endpoint = DEFAULT_PROTOCOL_ENDPOINT
+    _protocol_endpoint = KNOWN_WIRESERVER_IP
     _counts = {"hostplugin":0, "protocol":0, "other":0}
 
     @staticmethod
@@ -131,7 +131,7 @@ class IOErrorCounter(object):
             IOErrorCounter._counts = {"hostplugin":0, "protocol":0, "other":0}
 
     @staticmethod
-    def set_protocol_endpoint(endpoint=DEFAULT_PROTOCOL_ENDPOINT):
+    def set_protocol_endpoint(endpoint=KNOWN_WIRESERVER_IP):
         IOErrorCounter._protocol_endpoint = endpoint
 
 


Test patch:
diff --git a/tests/protocol/test_protocol_util.py b/tests/protocol/test_protocol_util.py
index 4a117899..ef6a2c6c 100644
--- a/tests/protocol/test_protocol_util.py
+++ b/tests/protocol/test_protocol_util.py
@@ -19,6 +19,9 @@ from tests.tools import *
 from azurelinuxagent.common.exception import *
 from azurelinuxagent.common.protocol import get_protocol_util, \
                                             TAG_FILE_NAME
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP
+from azurelinuxagent.common.protocol.util import ENDPOINT_FILE_NAME
+from errno import ENOENT
 
 @patch("time.sleep")
 class TestProtocolUtil(AgentTestCase):
@@ -34,22 +37,22 @@ class TestProtocolUtil(AgentTestCase):
         protocol_util.dhcp_handler = MagicMock()
         protocol_util.dhcp_handler.endpoint = "foo.bar"
 
-        #Test wire protocol is available
+        # Test wire protocol is available
         protocol = protocol_util.get_protocol()
         self.assertEquals(WireProtocol.return_value, protocol)
 
-        #Test wire protocol is not available
+        # Test wire protocol is not available
         protocol_util.clear_protocol()
         WireProtocol.return_value.detect.side_effect = ProtocolError()
 
         protocol = protocol_util.get_protocol()
         self.assertEquals(MetadataProtocol.return_value, protocol)
 
-        #Test no protocol is available
+        # Test no protocol is available
         protocol_util.clear_protocol()
         WireProtocol.return_value.detect.side_effect = ProtocolError()
-
         MetadataProtocol.return_value.detect.side_effect = ProtocolError()
+
         self.assertRaises(ProtocolError, protocol_util.get_protocol)
 
     def test_detect_protocol_by_file(self, _):
@@ -59,12 +62,12 @@ class TestProtocolUtil(AgentTestCase):
 
         tag_file = os.path.join(self.tmp_dir, TAG_FILE_NAME)
 
-        #Test tag file doesn't exist
+        # Test tag file doesn't exist
         protocol_util.get_protocol(by_file=True)
         protocol_util._detect_wire_protocol.assert_any_call()
         protocol_util._detect_metadata_protocol.assert_not_called()
 
-        #Test tag file exists
+        # Test tag file exists
         protocol_util.clear_protocol()
         protocol_util._detect_wire_protocol.reset_mock()
         protocol_util._detect_metadata_protocol.reset_mock()
@@ -75,6 +78,156 @@ class TestProtocolUtil(AgentTestCase):
         protocol_util._detect_metadata_protocol.assert_any_call()
         protocol_util._detect_wire_protocol.assert_not_called()
 
+    @patch("azurelinuxagent.common.conf.get_lib_dir")
+    @patch("azurelinuxagent.common.protocol.util.WireProtocol")
+    def test_detect_wire_protocol_no_dhcp(self, WireProtocol, mock_get_lib_dir, _):
+        WireProtocol.return_value.detect = Mock()
+        mock_get_lib_dir.return_value = self.tmp_dir
+
+        protocol_util = get_protocol_util()
+
+        protocol_util.osutil = MagicMock()
+        protocol_util.osutil.is_dhcp_available.return_value = False
+
+        protocol_util.dhcp_handler = MagicMock()
+        protocol_util.dhcp_handler.endpoint = None
+        protocol_util.dhcp_handler.run = Mock()
+
+        endpoint_file = protocol_util._get_wireserver_endpoint_file_path()
+
+        # Test wire protocol when no endpoint file has been written
+        protocol_util._detect_wire_protocol()
+        self.assertEqual(KNOWN_WIRESERVER_IP, protocol_util.get_wireserver_endpoint())
+
+        # Test wire protocol when endpoint was previously detected
+        protocol_util.clear_protocol()
+        with open(endpoint_file, "w+") as endpoint_fd:
+            endpoint_fd.write("baz.qux")
+
+        protocol_util._detect_wire_protocol()
+        self.assertEqual("baz.qux", protocol_util.get_wireserver_endpoint())
+
+        # Test wire protocol on dhcp failure
+        protocol_util.clear_protocol()
+        protocol_util.osutil.is_dhcp_available.return_value = True
+        protocol_util.dhcp_handler.run.side_effect = DhcpError()
+
+        self.assertRaises(ProtocolError, protocol_util._detect_wire_protocol)
+
+    @patch("azurelinuxagent.common.protocol.util.MetadataProtocol")
+    @patch("azurelinuxagent.common.protocol.util.WireProtocol")
+    def test_get_protocol(self, WireProtocol, MetadataProtocol, _):
+        WireProtocol.return_value = MagicMock()
+        MetadataProtocol.return_value = MagicMock()
+
+        protocol_util = get_protocol_util()
+        protocol_util.get_wireserver_endpoint = Mock()
+        protocol_util._detect_protocol = MagicMock()
+
+        # Test for wire protocol
+        protocol_util._save_protocol("WireProtocol")
+
+        protocol = protocol_util.get_protocol()
+        self.assertEquals(WireProtocol.return_value, protocol)
+        protocol_util.get_wireserver_endpoint.assert_any_call()
+
+        # Test to ensure protocol persists
+        protocol_util.get_wireserver_endpoint.reset_mock()
+        protocol_util._save_protocol("MetadataProtocol")
+
+        protocol = protocol_util.get_protocol()
+        self.assertEquals(WireProtocol.return_value, protocol)
+        protocol_util.get_wireserver_endpoint.assert_not_called()
+
+        # Test for metadata protocol
+        protocol_util.clear_protocol()
+        protocol_util._save_protocol("MetadataProtocol")
+
+        protocol = protocol_util.get_protocol()
+        self.assertEquals(MetadataProtocol.return_value, protocol)
+        protocol_util.get_wireserver_endpoint.assert_not_called()
+
+        # Test for unknown protocol
+        protocol_util.clear_protocol()
+        protocol_util._save_protocol("Not_a_Protocol")
+        protocol_util._detect_protocol.side_effect = NotImplementedError()
+
+        self.assertRaises(NotImplementedError, protocol_util.get_protocol)
+        protocol_util.get_wireserver_endpoint.assert_not_called()
+
+    @patch("azurelinuxagent.common.utils.fileutil")
+    @patch("azurelinuxagent.common.conf.get_lib_dir")
+    def test_endpoint_file_states(self, mock_get_lib_dir, mock_fileutil, _):
+        mock_get_lib_dir.return_value = self.tmp_dir
+        mock_fileutil = MagicMock()
+
+        protocol_util = get_protocol_util()
+        endpoint_file = protocol_util._get_wireserver_endpoint_file_path()
+
+        # Test get endpoint for io error
+        mock_fileutil.read_file.side_effect = IOError()
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertEquals(ep, KNOWN_WIRESERVER_IP)
+
+        # Test get endpoint when file not found
+        mock_fileutil.read_file.side_effect = IOError(ENOENT, 'File not found')
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertEquals(ep, KNOWN_WIRESERVER_IP)
+
+        # Test get endpoint for empty file
+        mock_fileutil.read_file.return_value = ""
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertEquals(ep, KNOWN_WIRESERVER_IP)
+
+        # Test set endpoint for io error
+        mock_fileutil.write_file.side_effect = IOError()
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertRaises(OSUtilError, protocol_util._set_wireserver_endpoint('abc'))
+
+        # Test clear endpoint for io error
+        with open(endpoint_file, "w+") as ep_fd:
+            ep_fd.write("")
+
+        with patch('os.remove') as mock_remove:
+            protocol_util._clear_wireserver_endpoint()
+            self.assertEqual(1, mock_remove.call_count)
+            self.assertEqual(endpoint_file, mock_remove.call_args_list[0][0][0])
+
+        # Test clear endpoint when file not found
+        with patch('os.remove') as mock_remove:
+            mock_remove = Mock(side_effect=IOError(ENOENT, 'File not found'))
+            protocol_util._clear_wireserver_endpoint()
+            mock_remove.assert_not_called()
+
+    def test_protocol_file_states(self, _):
+        protocol_util = get_protocol_util()
+        protocol_util._clear_wireserver_endpoint = Mock()
+
+        protocol_file = protocol_util._get_protocol_file_path()
+
+        # Test clear protocol for io error
+        with open(protocol_file, "w+") as proto_fd:
+            proto_fd.write("")
+
+        with patch('os.remove') as mock_remove:
+            protocol_util.clear_protocol()
+            self.assertEqual(1, protocol_util._clear_wireserver_endpoint.call_count)
+            self.assertEqual(1, mock_remove.call_count)
+            self.assertEqual(protocol_file, mock_remove.call_args_list[0][0][0])
+
+        # Test clear protocol when file not found
+        protocol_util._clear_wireserver_endpoint.reset_mock()
+
+        with patch('os.remove') as mock_remove:
+            protocol_util.clear_protocol()
+            self.assertEqual(1, protocol_util._clear_wireserver_endpoint.call_count)
+            self.assertEqual(1, mock_remove.call_count)
+            self.assertEqual(protocol_file, mock_remove.call_args_list[0][0][0])
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/utils/test_rest_util.py b/tests/utils/test_rest_util.py
index f097f10f..663f93a8 100644
--- a/tests/utils/test_rest_util.py
+++ b/tests/utils/test_rest_util.py
@@ -28,7 +28,7 @@ class TestIOErrorCounter(AgentTestCase):
         restutil.IOErrorCounter.set_protocol_endpoint()
 
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+            restutil.KNOWN_WIRESERVER_IP, restutil.HOST_PLUGIN_PORT)
 
         counts = restutil.IOErrorCounter.get_and_reset()
         self.assertEqual(1, counts["hostplugin"])
@@ -40,7 +40,7 @@ class TestIOErrorCounter(AgentTestCase):
         restutil.IOErrorCounter.set_protocol_endpoint()
 
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, 80)
+            restutil.KNOWN_WIRESERVER_IP, 80)
 
         counts = restutil.IOErrorCounter.get_and_reset()
         self.assertEqual(0, counts["hostplugin"])
@@ -64,11 +64,11 @@ class TestIOErrorCounter(AgentTestCase):
         restutil.IOErrorCounter.set_protocol_endpoint()
 
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+            restutil.KNOWN_WIRESERVER_IP, restutil.HOST_PLUGIN_PORT)
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+            restutil.KNOWN_WIRESERVER_IP, restutil.HOST_PLUGIN_PORT)
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, 80)
+            restutil.KNOWN_WIRESERVER_IP, 80)
         restutil.IOErrorCounter.increment(
             '169.254.169.254', 80)
         restutil.IOErrorCounter.increment(