From d590639a7891c4d746ca5f7a19a500d427009199 Mon Sep 17 00:00:00 2001
From: Sae126V <saitejav2021@gmail.com>
Date: Tue, 28 Mar 2023 14:34:34 +0000
Subject: [PATCH] Update exceptions to catch the timeout errors

---
 ssm/agents.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/ssm/agents.py b/ssm/agents.py
index 46e401d7..3d6ef885 100644
--- a/ssm/agents.py
+++ b/ssm/agents.py
@@ -37,7 +37,8 @@
 
 from stomp.exception import NotConnectedException
 try:
-    from argo_ams_library import AmsConnectionException
+    from argo_ams_library import (AmsConnectionException, AmsTimeoutException,
+                                  AmsBalancerException)
 except ImportError:
     # ImportError is raised when Ssm2 initialised if AMS is requested but lib
     # not installed.
@@ -328,7 +329,9 @@ def run_receiver(protocol, brokers, project, token, cp, log, dn_file):
                     if protocol == Ssm2.STOMP_MESSAGING:
                         ssm.send_ping()
 
-            except (NotConnectedException, AmsConnectionException) as error:
+            except (NotConnectedException, AmsConnectionException,
+                    AmsTimeoutException, AmsBalancerException) as error:
+
                 log.warning('Connection lost.')
                 log.debug(error)
                 ssm.shutdown()