Merge pull request #927 from Numerlor/truncate-internal-eval

Truncate internal eval.

7 files changed, 186 insertions, 62 deletions

diff --git a/bot/cogs/eval.py b/bot/cogs/eval.py
index eb8bfb1cf..23e5998d8 100644
--- a/bot/cogs/eval.py
+++ b/bot/cogs/eval.py
@@ -15,6 +15,7 @@ from bot.bot import Bot
 from bot.constants import Roles
 from bot.decorators import with_role
 from bot.interpreter import Interpreter
+from bot.utils import find_nth_occurrence, send_to_paste_service
 
 log = logging.getLogger(__name__)
 
@@ -171,6 +172,30 @@ async def func():  # (None,) -> Any
             res = traceback.format_exc()
 
         out, embed = self._format(code, res)
+        out = out.rstrip("\n")  # Strip empty lines from output
+
+        # Truncate output to max 15 lines or 1500 characters
+        newline_truncate_index = find_nth_occurrence(out, "\n", 15)
+
+        if newline_truncate_index is None or newline_truncate_index > 1500:
+            truncate_index = 1500
+        else:
+            truncate_index = newline_truncate_index
+
+        if len(out) > truncate_index:
+            paste_link = await send_to_paste_service(self.bot.http_session, out, extension="py")
+            if paste_link is not None:
+                paste_text = f"full contents at {paste_link}"
+            else:
+                paste_text = "failed to upload contents to paste service."
+
+            await ctx.send(
+                f"```py\n{out[:truncate_index]}\n```"
+                f"... response truncated; {paste_text}",
+                embed=embed
+            )
+            return
+
         await ctx.send(f"```py\n{out}```", embed=embed)
 
     @group(name='internal', aliases=('int',))
diff --git a/bot/cogs/snekbox.py b/bot/cogs/snekbox.py
index 63e6d7f31..03bf454ac 100644
--- a/bot/cogs/snekbox.py
+++ b/bot/cogs/snekbox.py
@@ -14,6 +14,7 @@ from discord.ext.commands import Cog, Context, command, guild_only
 from bot.bot import Bot
 from bot.constants import Categories, Channels, Roles, URLs
 from bot.decorators import in_whitelist
+from bot.utils import send_to_paste_service
 from bot.utils.messages import wait_for_deletion
 
 log = logging.getLogger(__name__)
@@ -71,17 +72,7 @@ class Snekbox(Cog):
         if len(output) > MAX_PASTE_LEN:
             log.info("Full output is too long to upload")
             return "too long to upload"
-
-        url = URLs.paste_service.format(key="documents")
-        try:
-            async with self.bot.http_session.post(url, data=output, raise_for_status=True) as resp:
-                data = await resp.json()
-
-            if "key" in data:
-                return URLs.paste_service.format(key=data["key"])
-        except Exception:
-            # 400 (Bad Request) means there are too many characters
-            log.exception("Failed to upload full output to paste service!")
+        return await send_to_paste_service(self.bot.http_session, output, extension="txt")
 
     @staticmethod
     def prepare_input(code: str) -> str:
diff --git a/bot/utils/__init__.py b/bot/utils/__init__.py
index 5a6e1811b..3e93fcb06 100644
--- a/bot/utils/__init__.py
+++ b/bot/utils/__init__.py
@@ -1,18 +1,5 @@
-from abc import ABCMeta
-
-from discord.ext.commands import CogMeta
-
+from bot.utils.helpers import CogABCMeta, find_nth_occurrence, pad_base64
 from bot.utils.redis_cache import RedisCache
+from bot.utils.services import send_to_paste_service
 
-__all__ = ['RedisCache', 'CogABCMeta']
-
-
-class CogABCMeta(CogMeta, ABCMeta):
-    """Metaclass for ABCs meant to be implemented as Cogs."""
-
-    pass
-
-
-def pad_base64(data: str) -> str:
-    """Return base64 `data` with padding characters to ensure its length is a multiple of 4."""
-    return data + "=" * (-len(data) % 4)
+__all__ = ['RedisCache', 'CogABCMeta', 'find_nth_occurrence', 'pad_base64', 'send_to_paste_service']
diff --git a/bot/utils/helpers.py b/bot/utils/helpers.py
new file mode 100644
index 000000000..d9b60af07
--- /dev/null
+++ b/bot/utils/helpers.py
@@ -0,0 +1,23 @@
+from abc import ABCMeta
+from typing import Optional
+
+from discord.ext.commands import CogMeta
+
+
+class CogABCMeta(CogMeta, ABCMeta):
+    """Metaclass for ABCs meant to be implemented as Cogs."""
+
+
+def find_nth_occurrence(string: str, substring: str, n: int) -> Optional[int]:
+    """Return index of `n`th occurrence of `substring` in `string`, or None if not found."""
+    index = 0
+    for _ in range(n):
+        index = string.find(substring, index+1)
+        if index == -1:
+            return None
+    return index
+
+
+def pad_base64(data: str) -> str:
+    """Return base64 `data` with padding characters to ensure its length is a multiple of 4."""
+    return data + "=" * (-len(data) % 4)
diff --git a/bot/utils/services.py b/bot/utils/services.py
new file mode 100644
index 000000000..087b9f969
--- /dev/null
+++ b/bot/utils/services.py
@@ -0,0 +1,54 @@
+import logging
+from typing import Optional
+
+from aiohttp import ClientConnectorError, ClientSession
+
+from bot.constants import URLs
+
+log = logging.getLogger(__name__)
+
+FAILED_REQUEST_ATTEMPTS = 3
+
+
+async def send_to_paste_service(http_session: ClientSession, contents: str, *, extension: str = "") -> Optional[str]:
+    """
+    Upload `contents` to the paste service.
+
+    `http_session` should be the current running ClientSession from aiohttp
+    `extension` is added to the output URL
+
+    When an error occurs, `None` is returned, otherwise the generated URL with the suffix.
+    """
+    extension = extension and f".{extension}"
+    log.debug(f"Sending contents of size {len(contents.encode())} bytes to paste service.")
+    paste_url = URLs.paste_service.format(key="documents")
+    for attempt in range(1, FAILED_REQUEST_ATTEMPTS + 1):
+        try:
+            async with http_session.post(paste_url, data=contents) as response:
+                response_json = await response.json()
+        except ClientConnectorError:
+            log.warning(
+                f"Failed to connect to paste service at url {paste_url}, "
+                f"trying again ({attempt}/{FAILED_REQUEST_ATTEMPTS})."
+            )
+            continue
+        except Exception:
+            log.exception(
+                f"An unexpected error has occurred during handling of the request, "
+                f"trying again ({attempt}/{FAILED_REQUEST_ATTEMPTS})."
+            )
+            continue
+
+        if "message" in response_json:
+            log.warning(
+                f"Paste service returned error {response_json['message']} with status code {response.status}, "
+                f"trying again ({attempt}/{FAILED_REQUEST_ATTEMPTS})."
+            )
+            continue
+        elif "key" in response_json:
+            log.info(f"Successfully uploaded contents to paste service behind key {response_json['key']}.")
+            return URLs.paste_service.format(key=response_json['key']) + extension
+        log.warning(
+            f"Got unexpected JSON response from paste service: {response_json}\n"
+            f"trying again ({attempt}/{FAILED_REQUEST_ATTEMPTS})."
+        )
diff --git a/tests/bot/cogs/test_snekbox.py b/tests/bot/cogs/test_snekbox.py
index 343e37db9..f22952931 100644
--- a/tests/bot/cogs/test_snekbox.py
+++ b/tests/bot/cogs/test_snekbox.py
@@ -1,5 +1,4 @@
 import asyncio
-import logging
 import unittest
 from unittest.mock import AsyncMock, MagicMock, Mock, call, create_autospec, patch
 
@@ -39,43 +38,14 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase):
         result = await self.cog.upload_output("-" * (snekbox.MAX_PASTE_LEN + 1))
         self.assertEqual(result, "too long to upload")
 
-    async def test_upload_output(self):
+    @patch("bot.cogs.snekbox.send_to_paste_service")
+    async def test_upload_output(self, mock_paste_util):
         """Upload the eval output to the URLs.paste_service.format(key="documents") endpoint."""
-        key = "MarkDiamond"
-        resp = MagicMock()
-        resp.json = AsyncMock(return_value={"key": key})
-
-        context_manager = MagicMock()
-        context_manager.__aenter__.return_value = resp
-        self.bot.http_session.post.return_value = context_manager
-
-        self.assertEqual(
-            await self.cog.upload_output("My awesome output"),
-            constants.URLs.paste_service.format(key=key)
-        )
-        self.bot.http_session.post.assert_called_with(
-            constants.URLs.paste_service.format(key="documents"),
-            data="My awesome output",
-            raise_for_status=True
+        await self.cog.upload_output("Test output.")
+        mock_paste_util.assert_called_once_with(
+            self.bot.http_session, "Test output.", extension="txt"
         )
 
-    async def test_upload_output_gracefully_fallback_if_exception_during_request(self):
-        """Output upload gracefully fallback if the upload fail."""
-        resp = MagicMock()
-        resp.json = AsyncMock(side_effect=Exception)
-
-        context_manager = MagicMock()
-        context_manager.__aenter__.return_value = resp
-        self.bot.http_session.post.return_value = context_manager
-
-        log = logging.getLogger("bot.cogs.snekbox")
-        with self.assertLogs(logger=log, level='ERROR'):
-            await self.cog.upload_output('My awesome output!')
-
-    async def test_upload_output_gracefully_fallback_if_no_key_in_response(self):
-        """Output upload gracefully fallback if there is no key entry in the response body."""
-        self.assertEqual((await self.cog.upload_output('My awesome output!')), None)
-
     def test_prepare_input(self):
         cases = (
             ('print("Hello world!")', 'print("Hello world!")', 'non-formatted'),
diff --git a/tests/bot/utils/test_services.py b/tests/bot/utils/test_services.py
new file mode 100644
index 000000000..5e0855704
--- /dev/null
+++ b/tests/bot/utils/test_services.py
@@ -0,0 +1,74 @@
+import logging
+import unittest
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+
+from aiohttp import ClientConnectorError
+
+from bot.utils.services import FAILED_REQUEST_ATTEMPTS, send_to_paste_service
+
+
+class PasteTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self) -> None:
+        self.http_session = MagicMock()
+
+    @patch("bot.utils.services.URLs.paste_service", "https://paste_service.com/{key}")
+    async def test_url_and_sent_contents(self):
+        """Correct url was used and post was called with expected data."""
+        response = MagicMock(
+            json=AsyncMock(return_value={"key": ""})
+        )
+        self.http_session.post().__aenter__.return_value = response
+        self.http_session.post.reset_mock()
+        await send_to_paste_service(self.http_session, "Content")
+        self.http_session.post.assert_called_once_with("https://paste_service.com/documents", data="Content")
+
+    @patch("bot.utils.services.URLs.paste_service", "https://paste_service.com/{key}")
+    async def test_paste_returns_correct_url_on_success(self):
+        """Url with specified extension is returned on successful requests."""
+        key = "paste_key"
+        test_cases = (
+            (f"https://paste_service.com/{key}.txt", "txt"),
+            (f"https://paste_service.com/{key}.py", "py"),
+            (f"https://paste_service.com/{key}", ""),
+        )
+        response = MagicMock(
+            json=AsyncMock(return_value={"key": key})
+        )
+        self.http_session.post().__aenter__.return_value = response
+
+        for expected_output, extension in test_cases:
+            with self.subTest(msg=f"Send contents with extension {repr(extension)}"):
+                self.assertEqual(
+                    await send_to_paste_service(self.http_session, "", extension=extension),
+                    expected_output
+                )
+
+    async def test_request_repeated_on_json_errors(self):
+        """Json with error message and invalid json are handled as errors and requests repeated."""
+        test_cases = ({"message": "error"}, {"unexpected_key": None}, {})
+        self.http_session.post().__aenter__.return_value = response = MagicMock()
+        self.http_session.post.reset_mock()
+
+        for error_json in test_cases:
+            with self.subTest(error_json=error_json):
+                response.json = AsyncMock(return_value=error_json)
+                result = await send_to_paste_service(self.http_session, "")
+                self.assertEqual(self.http_session.post.call_count, FAILED_REQUEST_ATTEMPTS)
+                self.assertIsNone(result)
+
+            self.http_session.post.reset_mock()
+
+    async def test_request_repeated_on_connection_errors(self):
+        """Requests are repeated in the case of connection errors."""
+        self.http_session.post = MagicMock(side_effect=ClientConnectorError(Mock(), Mock()))
+        result = await send_to_paste_service(self.http_session, "")
+        self.assertEqual(self.http_session.post.call_count, FAILED_REQUEST_ATTEMPTS)
+        self.assertIsNone(result)
+
+    async def test_general_error_handled_and_request_repeated(self):
+        """All `Exception`s are handled, logged and request repeated."""
+        self.http_session.post = MagicMock(side_effect=Exception)
+        result = await send_to_paste_service(self.http_session, "")
+        self.assertEqual(self.http_session.post.call_count, FAILED_REQUEST_ATTEMPTS)
+        self.assertLogs("bot.utils", logging.ERROR)
+        self.assertIsNone(result)