Move exception handling from send_msg_ams() to send_all

Also replace argo_id != "TOO LARGE" with continue to simply exit out of loop iteration.

Signed-off-by: jounaidr <[email protected]>

ssm/ssm2.py

@@ -16,7 +16,6 @@
 from __future__ import print_function
 
 import pkg_resources
-from pympler import asizeof
 
 from ssm import crypto
 from ssm.message_directory import MessageDirectory
@@ -38,7 +37,7 @@
 from logging import getLogger, INFO, WARNING, DEBUG
 
 try:
-    from argo_ams_library import ArgoMessagingService, AmsMessage
+    from argo_ams_library import ArgoMessagingService, AmsMessage, AmsServiceException
 except ImportError:
     # ImportError is raised later on if AMS is requested but lib not installed.
     ArgoMessagingService = None
@@ -391,15 +390,8 @@ def _send_msg_ams(self, text, msgid):
             message = AmsMessage(data=to_send,
                                  attributes={'empaid': msgid}).dict()
 
-            # Check the encrypted and wrapped message is less than 1mb
-            if asizeof.asizeof(message) < 1048576:
-                argo_response = self._ams.publish(self._dest, message, retry=3)
-                return argo_response['messageIds'][0]
-
-            else:
-                # Return a non AMS response so that the large message is not removed
-                log.warn('Message %s could not be sent as its larger than 1MB', msgid)
-                return "TOO LARGE"
+            argo_response = self._ams.publish(self._dest, message, retry=3)
+            return argo_response['messageIds'][0]
 
     def pull_msg_ams(self):
         """Pull 1 message from the AMS and acknowledge it."""
@@ -497,22 +489,30 @@ def send_all(self):
                 log_string = "Sent %s" % msgid
 
             elif self._protocol == Ssm2.AMS_MESSAGING:
-                # Then we are sending to an Argo Messaging Service instance.
-                argo_id = self._send_msg_ams(text, msgid)
-
-                log_string = "Sent %s, Argo ID: %s" % (msgid, argo_id)
+                try:
+                    # Then we are sending to an Argo Messaging Service instance.
+                    argo_id = self._send_msg_ams(text, msgid)
+                    log_string = "Sent %s, Argo ID: %s" % (msgid, argo_id)
+
+                except AmsServiceException as e:
+                    # Catch specific 'message too large' exception, raise otherwise.
+                    if "Message size is too large" not in str(e):
+                        raise
+                    else:
+                        # Exit out of loop iteration so that message is not removed.
+                        log.warn('Message %s could not be sent as its larger than 1MB', msgid)
+                        continue
 
             else:
                 # The SSM has been improperly configured
                 raise Ssm2Exception('Unknown messaging protocol: %s' %
                                     self._protocol)
 
-            if self._protocol == Ssm2.AMS_MESSAGING and argo_id != "TOO LARGE":
-                # log that the message was sent
-                log.info(log_string)
+            # log that the message was sent
+            log.info(log_string)
 
-                self._last_msg = None
-                self._outq.remove(msgid)
+            self._last_msg = None
+            self._outq.remove(msgid)
 
         log.info('Tidying message directory.')
         try: