completed receiveFromSQS.py

slack-bot/traffic-scraper/prd/receiveFromSQS.py

@@ -1,6 +1,8 @@
 import boto3
 import base64
 import math
+import json
+from datetime import datetime
 
 
 class receiveFromSQS(object):
@@ -35,20 +37,18 @@ class receiveFromSQS(object):
         self.loops = int(
             math.ceil(totalNumberOfMessages / maxNumberOfMessages)
         )
+        loopTrack = 0
         if totalNumberOfMessages <= 10:
             maxNumberOfMessages = totalNumberOfMessages
         else:
-            #  Find how many times total-loop*max, on last loop
-            #  replace max with this
-            maxNumberOfMessagesFinal = self.loops * maxNumberOfMessages - abs(10 - (
-                totalNumberOfMessages - (self.loops * maxNumberOfMessages)
-            ))
+            maxNumberOfMessagesFinal = 10 - (
+                (self.loops * maxNumberOfMessages) - totalNumberOfMessages
+            )
             print(maxNumberOfMessagesFinal)
         if self.loops == 0:
             raise RuntimeError('No messages in the queue')
         for i in range(0, self.loops):
-            loopTrack = 0
-            if loopTrack == self.loops:
+            if loopTrack == self.loops - 1 and totalNumberOfMessages > 10:
                 maxNumberOfMessages = maxNumberOfMessagesFinal
             self.resp.append(
                 self.sqs.receive_message(
@@ -56,12 +56,18 @@ class receiveFromSQS(object):
                     MaxNumberOfMessages=maxNumberOfMessages,
                 )
             )
-            entries = [  # Needs Keyerror try
-                {'Id': msg['MessageId'], 'ReceiptHandle': msg['ReceiptHandle']}
-                for msg in self.resp[i]['Messages']
-            ]
-            self._deleteSQSMessages(entries)
-            loopTrack += 1
+            try:
+                entries = [
+                    {
+                        'Id': msg['MessageId'],
+                        'ReceiptHandle': msg['ReceiptHandle'],
+                    }
+                    for msg in self.resp[i]['Messages']
+                ]
+                self._deleteSQSMessages(entries)
+                loopTrack += 1
+            except KeyError:
+                print("No messages in the queue")
         return self
 
     def _extractMessageFromSQS(self, totalNumberOfMessages):
@@ -119,7 +125,22 @@ class receiveFromSQS(object):
         self.receiveAllMessages(
             b64=b64, _totalNumberOfMessages=numberOfMessages
         )
-        pass
+        return self
+
+    def generateOutput(self, type='json'):
+        if type == 'json':
+            self.output = json.dumps(self.messages)
+        return self.output
+
+    def savetoDisk(self, path):
+        self.timeNow = datetime.now().strftime('%d-%m-%Y_%H:%M:%S')
+        if self.output is None:
+            self.generateOutput()
+        if len(self.messages) > 0:
+            with open(f'{path}/{self.timeNow}.json', 'w+') as outputFile:
+                outputFile.write(self.output)
+        else:
+            print('No messages to save')
 
 
 inst = receiveFromSQS.createSession(
@@ -128,6 +149,12 @@ inst = receiveFromSQS.createSession(
     '.com/745437999005/slack-bot.fifo',
 )
 
-inst.receiveNMessages(numberOfMessages=12)
+output = inst.receiveNMessages(numberOfMessages=6).generateOutput(type='json')
+inst.savetoDisk(
+    '/Users/dtomlinson/OneDrive - William Hill'
+    ' Organisation Limited/Mac/git_repos/python-VM/slack-bot/traffic-scraper'
+    '/prd'
+)
+# inst.receiveAllMessages()
 for item in inst.messages:
     print(item)