diff --git a/tests_e2e/tests/agent_publish/agent_publish.py b/tests_e2e/tests/agent_publish/agent_publish.py
index 397ecd0e3..eaddc74ed 100644
--- a/tests_e2e/tests/agent_publish/agent_publish.py
+++ b/tests_e2e/tests/agent_publish/agent_publish.py
@@ -57,13 +57,13 @@ def _get_agent_info(self) -> None:
 
     def _prepare_agent(self) -> None:
         log.info("Modifying agent update related config flags")
-        output = self._ssh_client.run_command("update-waagent-conf Debug.DownloadNewAgents=y AutoUpdate.GAFamily=Test", use_sudo=True)
-        log.info('Updated agent-update related config flags \n%s', output)
+        self._run_remote_test("update-waagent-conf Debug.DownloadNewAgents=y AutoUpdate.GAFamily=Test", use_sudo=True)
+        log.info('Updated agent-update DownloadNewAgents  GAFamily config flags')
 
     def _check_update(self) -> None:
         log.info("Verifying for agent update status")
-        output = self._ssh_client.run_command("agent_publish-check_update.py")
-        log.info('Checked the agent update \n%s', output)
+        self._run_remote_test("agent_publish-check_update.py")
+        log.info('Successfully checked the agent update')
 
     def _check_cse(self) -> None:
         custom_script_2_1 = VirtualMachineExtensionClient(
diff --git a/tests_e2e/tests/scripts/agent_publish-check_update.py b/tests_e2e/tests/scripts/agent_publish-check_update.py
index 9f8f66c4f..38ae00a90 100755
--- a/tests_e2e/tests/scripts/agent_publish-check_update.py
+++ b/tests_e2e/tests/scripts/agent_publish-check_update.py
@@ -17,10 +17,12 @@
 # limitations under the License.
 #
 import re
-import sys
-import logging
+
+from assertpy import fail
 
 from tests_e2e.tests.lib.agent_log import AgentLog
+from tests_e2e.tests.lib.logging import log
+from tests_e2e.tests.lib.remote_test import run_remote_test
 from tests_e2e.tests.lib.retry import retry_if_false
 
 
@@ -65,9 +67,9 @@ def verify_agent_update_from_log():
     update_successful = False
     update_version = ''
 
-    log = AgentLog()
+    agentlog = AgentLog()
 
-    for record in log.read():
+    for record in agentlog.read():
         if 'TelemetryData' in record.text:
             continue
 
@@ -76,37 +78,34 @@ def verify_agent_update_from_log():
             if update_match:
                 detected_update = True
                 update_version = update_match.groups()[2]
-                logging.info('found the agent update log: %s', record.text)
+                log.info('found the agent update log: %s', record.text)
                 break
 
         if detected_update:
             running_match = re.match(_RUNNING_PATTERN_00, record.text)
             if running_match and update_version == running_match.groups()[0]:
                 update_successful = True
-                logging.info('found the agent started new version log: %s', record.text)
+                log.info('found the agent started new version log: %s', record.text)
 
     if detected_update:
-        logging.info('update was detected: %s', update_version)
+        log.info('update was detected: %s', update_version)
         if update_successful:
-            logging.info('update was successful')
+            log.info('update was successful')
         else:
-            logging.warning('update was not successful')
+            log.warning('update was not successful')
             exit_code = 1
     else:
-        logging.warning('update was not detected')
+        log.warning('update was not detected')
         exit_code = 1
 
     return exit_code == 0
 
 
 # This method will trace agent update messages in the agent log and determine if the update was successful or not.
-try:
-    logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.DEBUG, stream=sys.stdout)
+def main():
     found: bool = retry_if_false(verify_agent_update_from_log)
     if not found:
-        raise Exception('update was not found in the logs')
-except Exception as e:
-    logging.error(e)
-    sys.exit(1)
+        fail('update was not found in the logs')
+
 
-sys.exit(0)
+run_remote_test(main)