[Snekbox] Fix up message consuming

2 files changed, 32 insertions, 16 deletions

diff --git a/bot/cogs/rmq.py b/bot/cogs/rmq.py
index 92645457e..c7063acd4 100644
--- a/bot/cogs/rmq.py
+++ b/bot/cogs/rmq.py
@@ -1,3 +1,4 @@
+import asyncio
 import datetime
 import json
 import logging
@@ -26,6 +27,8 @@ EMBED_PARAMS = (
     "colour", "title", "url", "description", "timestamp"
 )
 
+CONSUME_TIMEOUT = datetime.timedelta(seconds=10)
+
 
 class RMQ:
     """
@@ -65,7 +68,18 @@ class RMQ:
 
     async def consume(self, queue: str, **kwargs):
         queue_obj = await self.channel.declare_queue(queue, **kwargs)
-        return await queue_obj.get(timeout=30)
+
+        result = None
+        start_time = datetime.datetime.now()
+
+        while result is None:
+            if datetime.datetime.now() - start_time >= CONSUME_TIMEOUT:
+                result = "Timed out while waiting for a response."
+            else:
+                result = await queue_obj.get(timeout=5, fail=False)
+                await asyncio.sleep(0.5)
+
+        return result
 
     async def handle_message(self, message, data):
         log.debug(f"Message: {message}")
diff --git a/bot/cogs/snekbox.py b/bot/cogs/snekbox.py
index de979e46b..67bdd5a6e 100644
--- a/bot/cogs/snekbox.py
+++ b/bot/cogs/snekbox.py
@@ -52,21 +52,23 @@ class Snekbox:
                 snekid=str(ctx.author.id), message=code
             )
 
-            message = await self.rmq.consume(str(ctx.author.id), **RMQ_ARGS)
-            output = message.body.decode()
-
-            if "```" in output:
-                output = "Code block escape attempt detected; will not output result"
-            else:
-                output = [f"{i:03d} | {line}" for i, line in enumerate(output.split("\n"), start=1)]
-                output = "\n".join(output)
-
-                if len(output) >= 1900:
-                    output = f"{output[:1900]}... (truncated)"
-
-            await ctx.send(
-                f"{ctx.author.mention} Your eval job has completed.\n\n```{output}```"
-            )
+            async with ctx.typing():
+                message = await self.rmq.consume(str(ctx.author.id), **RMQ_ARGS)
+                output = message.body.decode().strip(" \n")  # Remove spaces and newlines from the ends
+
+                if "```" in output:
+                    output = "Code block escape attempt detected; will not output result"
+                else:
+                    if output.count("\n") > 0:
+                        output = [f"{i:03d} | {line}" for i, line in enumerate(output.split("\n"), start=1)]
+                        output = "\n".join(output)
+
+                    if len(output) >= 1900:
+                        output = f"{output[:1900]}... (truncated)"
+
+                await ctx.send(
+                    f"{ctx.author.mention} Your eval job has completed.\n\n```{output}```"
+                )
 
             del self.jobs[ctx.author.id]
         except Exception: