From 3f04611ddfc2e6d750d4c4e0a19d3cf154e7c5a9 Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Thu, 20 May 2021 16:09:39 -0400 Subject: chore: Update tests to correspond with the timeit command --- tests/bot/exts/utils/test_snekbox.py | 32 +++++++++++++++----------------- 1 file changed, 15 insertions(+), 17 deletions(-) (limited to 'tests') diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index 321a92445..1b3d61094 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -161,7 +161,9 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') self.cog.prepare_input.assert_called_once_with('MyAwesomeCode') - self.cog.send_eval.assert_called_once_with(ctx, 'MyAwesomeFormattedCode') + self.cog.send_eval.assert_called_once_with( + ctx, 'MyAwesomeFormattedCode', args=None, format_func=self.cog.format_output + ) self.cog.continue_eval.assert_called_once_with(ctx, response) async def test_eval_command_evaluate_twice(self): @@ -171,11 +173,13 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.prepare_input = MagicMock(return_value='MyAwesomeFormattedCode') self.cog.send_eval = AsyncMock(return_value=response) self.cog.continue_eval = AsyncMock() - self.cog.continue_eval.side_effect = ('MyAwesomeCode-2', None) + self.cog.continue_eval.side_effect = ('MyAwesomeFormattedCode', None) await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') self.cog.prepare_input.has_calls(call('MyAwesomeCode'), call('MyAwesomeCode-2')) - self.cog.send_eval.assert_called_with(ctx, 'MyAwesomeFormattedCode') + self.cog.send_eval.assert_called_with( + ctx, 'MyAwesomeFormattedCode', args=None, format_func=self.cog.format_output + ) self.cog.continue_eval.assert_called_with(ctx, response) async def test_eval_command_reject_two_eval_at_the_same_time(self): @@ -190,12 +194,6 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): "@LemonLemonishBeard#0042 You've already got a job running - please wait for it to finish!" ) - async def test_eval_command_call_help(self): - """Test if the eval command call the help command if no code is provided.""" - ctx = MockContext(command="sentinel") - await self.cog.eval_command(self.cog, ctx=ctx, code='') - ctx.send_help.assert_called_once_with(ctx.command) - async def test_send_eval(self): """Test the send_eval function.""" ctx = MockContext() @@ -212,11 +210,11 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_eval = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode') + await self.cog.send_eval(ctx, 'MyAwesomeCode', format_func=self.cog.format_output) ctx.send.assert_called_once_with( '@LemonLemonishBeard#0042 :yay!: Return code 0.\n\n```\n[No output]\n```' ) - self.cog.post_eval.assert_called_once_with('MyAwesomeCode') + self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': '', 'returncode': 0}) self.cog.get_results_message.assert_called_once_with({'stdout': '', 'returncode': 0}) self.cog.format_output.assert_called_once_with('') @@ -237,12 +235,12 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_eval = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode') + await self.cog.send_eval(ctx, 'MyAwesomeCode', format_func=self.cog.format_output) ctx.send.assert_called_once_with( '@LemonLemonishBeard#0042 :yay!: Return code 0.' '\n\n```\nWay too long beard\n```\nFull output: lookatmybeard.com' ) - self.cog.post_eval.assert_called_once_with('MyAwesomeCode') + self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': 'Way too long beard', 'returncode': 0}) self.cog.get_results_message.assert_called_once_with({'stdout': 'Way too long beard', 'returncode': 0}) self.cog.format_output.assert_called_once_with('Way too long beard') @@ -262,11 +260,11 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_eval = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode') + await self.cog.send_eval(ctx, 'MyAwesomeCode', format_func=self.cog.format_output) ctx.send.assert_called_once_with( '@LemonLemonishBeard#0042 :nope!: Return code 127.\n\n```\nBeard got stuck in the eval\n```' ) - self.cog.post_eval.assert_called_once_with('MyAwesomeCode') + self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': 'ERROR', 'returncode': 127}) self.cog.get_results_message.assert_called_once_with({'stdout': 'ERROR', 'returncode': 127}) self.cog.format_output.assert_not_called() @@ -282,7 +280,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.get_code = create_autospec(self.cog.get_code, spec_set=True, return_value=expected) actual = await self.cog.continue_eval(ctx, response) - self.cog.get_code.assert_awaited_once_with(new_msg) + self.cog.get_code.assert_awaited_once_with(new_msg, ctx.command) self.assertEqual(actual, expected) self.bot.wait_for.assert_has_awaits( ( @@ -327,7 +325,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.bot.get_context.return_value = MockContext(command=command) message = MockMessage(content=content) - actual_code = await self.cog.get_code(message) + actual_code = await self.cog.get_code(message, self.cog.eval_command) self.bot.get_context.assert_awaited_once_with(message) self.assertEqual(actual_code, expected_code) -- cgit v1.2.3 From b7e49a5fb1adb541db2cf5632a460a37ddda6d0a Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Thu, 13 Jan 2022 21:26:58 -0500 Subject: chore: Suppress output in the setup code, not the code that gets timed. If multiple formatted codeblocks are passed to the command, the first one will be used as the setup code that does not get timed. --- bot/exts/utils/snekbox.py | 70 +++++++++++++++++++++++++++--------- tests/bot/exts/utils/test_snekbox.py | 6 ++-- 2 files changed, 56 insertions(+), 20 deletions(-) (limited to 'tests') diff --git a/bot/exts/utils/snekbox.py b/bot/exts/utils/snekbox.py index bd521a4ee..0d8da5e56 100644 --- a/bot/exts/utils/snekbox.py +++ b/bot/exts/utils/snekbox.py @@ -2,9 +2,9 @@ import asyncio import contextlib import datetime import re -import textwrap from functools import partial from signal import Signals +from textwrap import dedent from typing import Awaitable, Callable, Optional, Tuple from discord import AllowedMentions, HTTPException, Message, NotFound, Reaction, User @@ -36,13 +36,35 @@ RAW_CODE_REGEX = re.compile( re.DOTALL # "." also matches newlines ) -TIMEIT_EVAL_WRAPPER = """ -from contextlib import redirect_stdout -from io import StringIO +TIMEIT_SETUP_WRAPPER = """ +import atexit +import sys +from collections import deque -with redirect_stdout(StringIO()): - del redirect_stdout, StringIO -{code} +if not hasattr(sys, "_setup_finished"): + class Writer(deque): + def __init__(self): + super().__init__(maxlen=1) + + def write(self, string): + if string.strip(): + self.append(string) + + def read(self): + return self.pop() + + def flush(self): + pass + + sys.stdout = Writer() + + def print_last_line(): + if sys.stdout: + print(sys.stdout.read(), file=sys.__stdout__) + + atexit.register(print_last_line) + sys._setup_finished = None +{setup} """ TIMEIT_OUTPUT_REGEX = re.compile(r"\d+ loops, best of \d+: \d(?:\.\d\d?)? [mnu]?sec per loop") @@ -90,34 +112,37 @@ class Snekbox(Cog): return await send_to_paste_service(output, extension="txt") @staticmethod - def prepare_input(code: str) -> str: + def prepare_input(code: str) -> list[str]: """ Extract code from the Markdown, format it, and insert it into the code template. If there is any code block, ignore text outside the code block. Use the first code block, but prefer a fenced code block. If there are several fenced code blocks, concatenate only the fenced code blocks. + + Retrun a list of code blocks if any, otherwise return a list with a single string of code. """ if match := list(FORMATTED_CODE_REGEX.finditer(code)): blocks = [block for block in match if block.group("block")] if len(blocks) > 1: - code = '\n'.join(block.group("code") for block in blocks) + codeblocks = [block.group("code") for block in blocks] info = "several code blocks" else: match = match[0] if len(blocks) == 0 else blocks[0] code, block, lang, delim = match.group("code", "block", "lang", "delim") + codeblocks = [dedent(code)] if block: info = (f"'{lang}' highlighted" if lang else "plain") + " code block" else: info = f"{delim}-enclosed inline code" else: - code = RAW_CODE_REGEX.fullmatch(code).group("code") + codeblocks = [dedent(RAW_CODE_REGEX.fullmatch(code).group("code"))] info = "unformatted or badly formatted code" - code = textwrap.dedent(code) + code = "\n".join(codeblocks) log.trace(f"Extracted {info} for evaluation:\n{code}") - return code + return codeblocks @staticmethod def get_results_message(results: dict) -> Tuple[str, str]: @@ -248,7 +273,7 @@ class Snekbox(Cog): log.info(f"{ctx.author}'s job had a return code of {results['returncode']}") return response - async def continue_eval(self, ctx: Context, response: Message) -> Optional[str]: + async def continue_eval(self, ctx: Context, response: Message) -> Optional[list[str]]: """ Check if the eval session should continue. @@ -380,7 +405,7 @@ class Snekbox(Cog): We've done our best to make this sandboxed, but do let us know if you manage to find an issue with it! """ - code = self.prepare_input(code) + code = "\n".join(self.prepare_input(code)) await self.run_eval(ctx, code, format_func=self.format_output) @command(name="timeit", aliases=("ti",)) @@ -400,13 +425,24 @@ class Snekbox(Cog): block. Code can be re-evaluated by editing the original message within 10 seconds and clicking the reaction that subsequently appears. + If multiple formatted codeblocks are provided, the first one will be the setup code, which will + not be timed. The remaining codeblocks will be joined together and timed. + We've done our best to make this sandboxed, but do let us know if you manage to find an issue with it! """ - code = self.prepare_input(code) + args = ["-m", "timeit"] + setup = "" + codeblocks = self.prepare_input(code) + + if len(codeblocks) > 1: + setup = codeblocks.pop(0) + + code = "\n".join(codeblocks) + args.extend(["-s", TIMEIT_SETUP_WRAPPER.format(setup=setup)]) + await self.run_eval( - ctx, TIMEIT_EVAL_WRAPPER.format(code=textwrap.indent(code, " ")), - format_func=self.format_timeit_output, args=["-m", "timeit"] + ctx, code=code, format_func=self.format_timeit_output, args=args ) diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index cbffaa6b0..ebab71e71 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -61,7 +61,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): ) for case, expected, testname in cases: with self.subTest(msg=f'Extract code from {testname}.'): - self.assertEqual(self.cog.prepare_input(case), expected) + self.assertEqual('\n'.join(self.cog.prepare_input(case)), expected) def test_get_results_message(self): """Return error and message according to the eval result.""" @@ -156,7 +156,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): """Test the eval command procedure.""" ctx = MockContext() response = MockMessage() - self.cog.prepare_input = MagicMock(return_value='MyAwesomeFormattedCode') + self.cog.prepare_input = MagicMock(return_value=['MyAwesomeFormattedCode']) self.cog.send_eval = AsyncMock(return_value=response) self.cog.continue_eval = AsyncMock(return_value=None) @@ -297,7 +297,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): actual = await self.cog.continue_eval(ctx, response) self.cog.get_code.assert_awaited_once_with(new_msg, ctx.command) - self.assertEqual(actual, expected) + self.assertEqual(actual, [expected]) self.bot.wait_for.assert_has_awaits( ( call( -- cgit v1.2.3 From 8594a3535413e662ae519f31989459a953e8a726 Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Mon, 17 Jan 2022 14:43:58 -0500 Subject: fix: Modify tests to correspond with Snekbox.continue_eval --- tests/bot/exts/utils/test_snekbox.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) (limited to 'tests') diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index ebab71e71..4245de8a3 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -156,32 +156,35 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): """Test the eval command procedure.""" ctx = MockContext() response = MockMessage() + ctx.command = MagicMock() + self.cog.prepare_input = MagicMock(return_value=['MyAwesomeFormattedCode']) self.cog.send_eval = AsyncMock(return_value=response) - self.cog.continue_eval = AsyncMock(return_value=None) + self.cog.continue_eval = AsyncMock(return_value=(None, None)) await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') self.cog.prepare_input.assert_called_once_with('MyAwesomeCode') self.cog.send_eval.assert_called_once_with( ctx, 'MyAwesomeFormattedCode', args=None, format_func=self.cog.format_output ) - self.cog.continue_eval.assert_called_once_with(ctx, response) + self.cog.continue_eval.assert_called_once_with(ctx, response, ctx.command) async def test_eval_command_evaluate_twice(self): """Test the eval and re-eval command procedure.""" ctx = MockContext() response = MockMessage() + ctx.command = MagicMock() self.cog.prepare_input = MagicMock(return_value='MyAwesomeFormattedCode') self.cog.send_eval = AsyncMock(return_value=response) self.cog.continue_eval = AsyncMock() - self.cog.continue_eval.side_effect = ('MyAwesomeFormattedCode', None) + self.cog.continue_eval.side_effect = (('MyAwesomeFormattedCode', None), (None, None)) await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') self.cog.prepare_input.has_calls(call('MyAwesomeCode'), call('MyAwesomeCode-2')) self.cog.send_eval.assert_called_with( ctx, 'MyAwesomeFormattedCode', args=None, format_func=self.cog.format_output ) - self.cog.continue_eval.assert_called_with(ctx, response) + self.cog.continue_eval.assert_called_with(ctx, response, ctx.command) async def test_eval_command_reject_two_eval_at_the_same_time(self): """Test if the eval command rejects an eval if the author already have a running eval.""" @@ -295,9 +298,9 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): expected = "NewCode" self.cog.get_code = create_autospec(self.cog.get_code, spec_set=True, return_value=expected) - actual = await self.cog.continue_eval(ctx, response) + actual = await self.cog.continue_eval(ctx, response, self.cog.eval_command) self.cog.get_code.assert_awaited_once_with(new_msg, ctx.command) - self.assertEqual(actual, [expected]) + self.assertEqual(actual, (expected, None)) self.bot.wait_for.assert_has_awaits( ( call( @@ -316,8 +319,8 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): ctx = MockContext(message=MockMessage(clear_reactions=AsyncMock())) self.bot.wait_for.side_effect = asyncio.TimeoutError - actual = await self.cog.continue_eval(ctx, MockMessage()) - self.assertEqual(actual, None) + actual = await self.cog.continue_eval(ctx, MockMessage(), self.cog.eval_command) + self.assertEqual(actual, (None, None)) ctx.message.clear_reaction.assert_called_once_with(snekbox.REEVAL_EMOJI) async def test_get_code(self): -- cgit v1.2.3 From 54e4f3777372ef526667885f4392030bab1b5b07 Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Mon, 17 Jan 2022 17:42:37 -0500 Subject: chore: Apply suggestions and adjust tests --- bot/exts/utils/snekbox.py | 52 +++++++++++++----------------------- tests/bot/exts/utils/test_snekbox.py | 24 ++++++++--------- 2 files changed, 29 insertions(+), 47 deletions(-) (limited to 'tests') diff --git a/bot/exts/utils/snekbox.py b/bot/exts/utils/snekbox.py index 49f1be17b..1d9646113 100644 --- a/bot/exts/utils/snekbox.py +++ b/bot/exts/utils/snekbox.py @@ -36,6 +36,8 @@ RAW_CODE_REGEX = re.compile( re.DOTALL # "." also matches newlines ) +# The timeit command should only output the very last line, so all other output should be suppressed. +# This will be used as the setup code along with any setup code provided. TIMEIT_SETUP_WRAPPER = """ import atexit import sys @@ -163,19 +165,19 @@ class Snekbox(Cog): return code, args @staticmethod - def get_results_message(results: dict) -> Tuple[str, str]: + def get_results_message(results: dict, job_name: str) -> Tuple[str, str]: """Return a user-friendly message and error corresponding to the process's return code.""" stdout, returncode = results["stdout"], results["returncode"] - msg = f"Your eval job has completed with return code {returncode}" + msg = f"Your {job_name} job has completed with return code {returncode}" error = "" if returncode is None: - msg = "Your eval job has failed" + msg = f"Your {job_name} job has failed" error = stdout.strip() elif returncode == 128 + SIGKILL: - msg = "Your eval job timed out or ran out of memory" + msg = f"Your {job_name} job timed out or ran out of memory" elif returncode == 255: - msg = "Your eval job has failed" + msg = f"Your {job_name} job has failed" error = "A fatal NsJail error occurred" else: # Try to append signal's name if one exists @@ -249,7 +251,7 @@ class Snekbox(Cog): code: str, *, args: Optional[list[str]] = None, - format_func: FormatFunc + job_name: str ) -> Message: """ Evaluate code, format it, and send the output to the corresponding channel. @@ -258,13 +260,13 @@ class Snekbox(Cog): """ async with ctx.typing(): results = await self.post_eval(code, args=args) - msg, error = self.get_results_message(results) + msg, error = self.get_results_message(results, job_name) if error: output, paste_link = error, None else: log.trace("Formatting output...") - output, paste_link = await format_func(results["stdout"]) + output, paste_link = await self.format_output(results["stdout"]) icon = self.get_status_emoji(results) msg = f"{ctx.author.mention} {icon} {msg}.\n\n```\n{output}\n```" @@ -288,12 +290,12 @@ class Snekbox(Cog): response = await ctx.send(msg, allowed_mentions=allowed_mentions) scheduling.create_task(wait_for_deletion(response, (ctx.author.id,)), event_loop=self.bot.loop) - log.info(f"{ctx.author}'s job had a return code of {results['returncode']}") + log.info(f"{ctx.author}'s {job_name} job had a return code of {results['returncode']}") return response async def continue_eval( self, ctx: Context, response: Message, command: Command - ) -> Optional[tuple[str, Optional[list[str]]]]: + ) -> tuple[Optional[str], Optional[list[str]]]: """ Check if the eval session should continue. @@ -355,19 +357,15 @@ class Snekbox(Cog): return code - async def run_eval( + async def run_job( self, + job_name: str, ctx: Context, code: str, - format_func: FormatFunc, *, args: Optional[list[str]] = None, ) -> None: - """ - Handles checks, stats and re-evaluation of an eval. - - `format_func` is an async callable that takes a string (the output) and formats it to show to the user. - """ + """Handles checks, stats and re-evaluation of a snekbox job.""" if ctx.author.id in self.jobs: await ctx.send( f"{ctx.author.mention} You've already got a job running - " @@ -392,7 +390,7 @@ class Snekbox(Cog): while True: self.jobs[ctx.author.id] = datetime.datetime.now() try: - response = await self.send_eval(ctx, code, args=args, format_func=format_func) + response = await self.send_eval(ctx, code, args=args, job_name=job_name) finally: del self.jobs[ctx.author.id] @@ -401,18 +399,6 @@ class Snekbox(Cog): break log.info(f"Re-evaluating code from message {ctx.message.id}:\n{code}") - async def format_timeit_output(self, output: str) -> tuple[str, str]: - """ - Parses the time from the end of the output given by timeit. - - If an error happened, then it won't contain the time and instead proceed with regular formatting. - """ - split_output = output.rstrip("\n").rsplit("\n", 1) - if len(split_output) == 2 and TIMEIT_OUTPUT_REGEX.fullmatch(split_output[1]): - return split_output[1], None - - return await self.format_output(output) - @command(name="eval", aliases=("e",)) @guild_only() @redirect_output( @@ -434,7 +420,7 @@ class Snekbox(Cog): issue with it! """ code = "\n".join(self.prepare_input(code)) - await self.run_eval(ctx, code, format_func=self.format_output) + await self.run_job("eval", ctx, code) @command(name="timeit", aliases=("ti",)) @guild_only() @@ -462,9 +448,7 @@ class Snekbox(Cog): codeblocks = self.prepare_input(code) code, args = self.prepare_timeit_input(codeblocks) - await self.run_eval( - ctx, code=code, format_func=self.format_timeit_output, args=args - ) + await self.run_job("timeit", ctx, code=code, args=args) def predicate_eval_message_edit(ctx: Context, old_msg: Message, new_msg: Message) -> bool: diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index 4245de8a3..339cdaaa4 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -72,13 +72,13 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): ) for stdout, returncode, expected in cases: with self.subTest(stdout=stdout, returncode=returncode, expected=expected): - actual = self.cog.get_results_message({'stdout': stdout, 'returncode': returncode}) + actual = self.cog.get_results_message({'stdout': stdout, 'returncode': returncode}, 'eval') self.assertEqual(actual, expected) @patch('bot.exts.utils.snekbox.Signals', side_effect=ValueError) def test_get_results_message_invalid_signal(self, mock_signals: Mock): self.assertEqual( - self.cog.get_results_message({'stdout': '', 'returncode': 127}), + self.cog.get_results_message({'stdout': '', 'returncode': 127}, 'eval'), ('Your eval job has completed with return code 127', '') ) @@ -86,7 +86,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): def test_get_results_message_valid_signal(self, mock_signals: Mock): mock_signals.return_value.name = 'SIGTEST' self.assertEqual( - self.cog.get_results_message({'stdout': '', 'returncode': 127}), + self.cog.get_results_message({'stdout': '', 'returncode': 127}, 'eval'), ('Your eval job has completed with return code 127 (SIGTEST)', '') ) @@ -164,9 +164,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') self.cog.prepare_input.assert_called_once_with('MyAwesomeCode') - self.cog.send_eval.assert_called_once_with( - ctx, 'MyAwesomeFormattedCode', args=None, format_func=self.cog.format_output - ) + self.cog.send_eval.assert_called_once_with(ctx, 'MyAwesomeFormattedCode', args=None, job_name='eval') self.cog.continue_eval.assert_called_once_with(ctx, response, ctx.command) async def test_eval_command_evaluate_twice(self): @@ -182,7 +180,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') self.cog.prepare_input.has_calls(call('MyAwesomeCode'), call('MyAwesomeCode-2')) self.cog.send_eval.assert_called_with( - ctx, 'MyAwesomeFormattedCode', args=None, format_func=self.cog.format_output + ctx, 'MyAwesomeFormattedCode', args=None, job_name='eval' ) self.cog.continue_eval.assert_called_with(ctx, response, ctx.command) @@ -214,7 +212,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_eval = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode', format_func=self.cog.format_output) + await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') ctx.send.assert_called_once() self.assertEqual( @@ -227,7 +225,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': '', 'returncode': 0}) - self.cog.get_results_message.assert_called_once_with({'stdout': '', 'returncode': 0}) + self.cog.get_results_message.assert_called_once_with({'stdout': '', 'returncode': 0}, 'eval') self.cog.format_output.assert_called_once_with('') async def test_send_eval_with_paste_link(self): @@ -246,7 +244,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_eval = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode', format_func=self.cog.format_output) + await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') ctx.send.assert_called_once() self.assertEqual( @@ -257,7 +255,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': 'Way too long beard', 'returncode': 0}) - self.cog.get_results_message.assert_called_once_with({'stdout': 'Way too long beard', 'returncode': 0}) + self.cog.get_results_message.assert_called_once_with({'stdout': 'Way too long beard', 'returncode': 0}, 'eval') self.cog.format_output.assert_called_once_with('Way too long beard') async def test_send_eval_with_non_zero_eval(self): @@ -275,7 +273,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_eval = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode', format_func=self.cog.format_output) + await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') ctx.send.assert_called_once() self.assertEqual( @@ -285,7 +283,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': 'ERROR', 'returncode': 127}) - self.cog.get_results_message.assert_called_once_with({'stdout': 'ERROR', 'returncode': 127}) + self.cog.get_results_message.assert_called_once_with({'stdout': 'ERROR', 'returncode': 127}, 'eval') self.cog.format_output.assert_not_called() @patch("bot.exts.utils.snekbox.partial") -- cgit v1.2.3 From d0cf7f2f1573e883cf1f6aaf8c54b3701496722b Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Wed, 26 Jan 2022 16:12:15 -0500 Subject: chore: Remove the naming of 'eval' in certain places Since the !eval command is no longer the only snekbox command, make the naming more generic. --- bot/exts/filters/filtering.py | 4 +- bot/exts/utils/snekbox.py | 72 ++++++++++++++++++------------------ tests/bot/exts/utils/test_snekbox.py | 6 +-- 3 files changed, 41 insertions(+), 41 deletions(-) (limited to 'tests') diff --git a/bot/exts/filters/filtering.py b/bot/exts/filters/filtering.py index ad904d147..e49cf4f82 100644 --- a/bot/exts/filters/filtering.py +++ b/bot/exts/filters/filtering.py @@ -267,9 +267,9 @@ class Filtering(Cog): # Update time when alert sent await self.name_alerts.set(member.id, arrow.utcnow().timestamp()) - async def filter_eval(self, result: str, msg: Message) -> bool: + async def filter_snekbox_job(self, result: str, msg: Message) -> bool: """ - Filter the result of an !eval to see if it violates any of our rules, and then respond accordingly. + Filter the result of a snekbox command to see if it violates any of our rules, and then respond accordingly. Also requires the original message, to check whether to filter and for mod logs. Returns whether a filter was triggered or not. diff --git a/bot/exts/utils/snekbox.py b/bot/exts/utils/snekbox.py index e7712eee5..86993b7f1 100644 --- a/bot/exts/utils/snekbox.py +++ b/bot/exts/utils/snekbox.py @@ -71,15 +71,15 @@ if not hasattr(sys, "_setup_finished"): MAX_PASTE_LEN = 10000 -# `!eval` command whitelists and blacklists. -NO_EVAL_CHANNELS = (Channels.python_general,) -NO_EVAL_CATEGORIES = () -EVAL_ROLES = (Roles.helpers, Roles.moderators, Roles.admins, Roles.owners, Roles.python_community, Roles.partners) +# The Snekbox commands' whitelists and blacklists. +NO_SNEKBOX_CHANNELS = (Channels.python_general,) +NO_SNEKBOX_CATEGORIES = () +SNEKBOX_ROLES = (Roles.helpers, Roles.moderators, Roles.admins, Roles.owners, Roles.python_community, Roles.partners) SIGKILL = 9 -REEVAL_EMOJI = '\U0001f501' # :repeat: -REEVAL_TIMEOUT = 30 +REDO_EMOJI = '\U0001f501' # :repeat: +REDO_TIMEOUT = 30 class Snekbox(Cog): @@ -89,7 +89,7 @@ class Snekbox(Cog): self.bot = bot self.jobs = {} - async def post_eval(self, code: str, *, args: Optional[list[str]] = None) -> dict: + async def post_job(self, code: str, *, args: Optional[list[str]] = None) -> dict: """Send a POST request to the Snekbox API to evaluate code and return the results.""" url = URLs.snekbox_eval_api data = {"input": code} @@ -101,7 +101,7 @@ class Snekbox(Cog): return await resp.json() async def upload_output(self, output: str) -> Optional[str]: - """Upload the eval output to a paste service and return a URL to it if successful.""" + """Upload the job's output to a paste service and return a URL to it if successful.""" log.trace("Uploading full output to paste service...") if len(output) > MAX_PASTE_LEN: @@ -241,7 +241,7 @@ class Snekbox(Cog): return output, paste_link - async def send_eval( + async def send_job( self, ctx: Context, code: str, @@ -255,7 +255,7 @@ class Snekbox(Cog): Return the bot response. """ async with ctx.typing(): - results = await self.post_eval(code, args=args) + results = await self.post_job(code, args=args) msg, error = self.get_results_message(results, job_name) if error: @@ -269,7 +269,7 @@ class Snekbox(Cog): if paste_link: msg = f"{msg}\nFull output: {paste_link}" - # Collect stats of eval fails + successes + # Collect stats of job fails + successes if icon == ":x:": self.bot.stats.incr("snekbox.python.fail") else: @@ -278,7 +278,7 @@ class Snekbox(Cog): filter_cog = self.bot.get_cog("Filtering") filter_triggered = False if filter_cog: - filter_triggered = await filter_cog.filter_eval(msg, ctx.message) + filter_triggered = await filter_cog.filter_snekbox_job(msg, ctx.message) if filter_triggered: response = await ctx.send("Attempt to circumvent filter detected. Moderator team has been alerted.") else: @@ -289,26 +289,26 @@ class Snekbox(Cog): log.info(f"{ctx.author}'s {job_name} job had a return code of {results['returncode']}") return response - async def continue_eval( + async def continue_job( self, ctx: Context, response: Message, command: Command ) -> tuple[Optional[str], Optional[list[str]]]: """ - Check if the eval session should continue. + Check if the job's session should continue. If the code is to be re-evaluated, return the new code, and the args if the command is the timeit command. - Otherwise return (None, None) if the eval session should be terminated. + Otherwise return (None, None) if the job's session should be terminated. """ - _predicate_eval_message_edit = partial(predicate_eval_message_edit, ctx) - _predicate_emoji_reaction = partial(predicate_eval_emoji_reaction, ctx) + _predicate_message_edit = partial(predicate_message_edit, ctx) + _predicate_emoji_reaction = partial(predicate_emoji_reaction, ctx) with contextlib.suppress(NotFound): try: _, new_message = await self.bot.wait_for( 'message_edit', - check=_predicate_eval_message_edit, - timeout=REEVAL_TIMEOUT + check=_predicate_message_edit, + timeout=REDO_TIMEOUT ) - await ctx.message.add_reaction(REEVAL_EMOJI) + await ctx.message.add_reaction(REDO_EMOJI) await self.bot.wait_for( 'reaction_add', check=_predicate_emoji_reaction, @@ -316,7 +316,7 @@ class Snekbox(Cog): ) code = await self.get_code(new_message, ctx.command) - await ctx.message.clear_reaction(REEVAL_EMOJI) + await ctx.message.clear_reaction(REDO_EMOJI) with contextlib.suppress(HTTPException): await response.delete() @@ -324,7 +324,7 @@ class Snekbox(Cog): return None, None except asyncio.TimeoutError: - await ctx.message.clear_reaction(REEVAL_EMOJI) + await ctx.message.clear_reaction(REDO_EMOJI) return None, None codeblocks = self.prepare_input(code) @@ -347,11 +347,11 @@ class Snekbox(Cog): new_ctx = await self.bot.get_context(message) if new_ctx.command is command: - log.trace(f"Message {message.id} invokes eval command.") + log.trace(f"Message {message.id} invokes {command} command.") split = message.content.split(maxsplit=1) code = split[1] if len(split) > 1 else None else: - log.trace(f"Message {message.id} does not invoke eval command.") + log.trace(f"Message {message.id} does not invoke {command} command.") code = message.content return code @@ -389,11 +389,11 @@ class Snekbox(Cog): while True: self.jobs[ctx.author.id] = datetime.datetime.now() try: - response = await self.send_eval(ctx, code, args=args, job_name=job_name) + response = await self.send_job(ctx, code, args=args, job_name=job_name) finally: del self.jobs[ctx.author.id] - code, args = await self.continue_eval(ctx, response, ctx.command) + code, args = await self.continue_job(ctx, response, ctx.command) if not code: break log.info(f"Re-evaluating code from message {ctx.message.id}:\n{code}") @@ -402,9 +402,9 @@ class Snekbox(Cog): @guild_only() @redirect_output( destination_channel=Channels.bot_commands, - bypass_roles=EVAL_ROLES, - categories=NO_EVAL_CATEGORIES, - channels=NO_EVAL_CHANNELS, + bypass_roles=SNEKBOX_ROLES, + categories=NO_SNEKBOX_CATEGORIES, + channels=NO_SNEKBOX_CHANNELS, ping_user=False ) async def eval_command(self, ctx: Context, *, code: str) -> None: @@ -425,9 +425,9 @@ class Snekbox(Cog): @guild_only() @redirect_output( destination_channel=Channels.bot_commands, - bypass_roles=EVAL_ROLES, - categories=NO_EVAL_CATEGORIES, - channels=NO_EVAL_CHANNELS, + bypass_roles=SNEKBOX_ROLES, + categories=NO_SNEKBOX_CATEGORIES, + channels=NO_SNEKBOX_CHANNELS, ping_user=False ) async def timeit_command(self, ctx: Context, *, code: str) -> str: @@ -450,14 +450,14 @@ class Snekbox(Cog): await self.run_job("timeit", ctx, code=code, args=args) -def predicate_eval_message_edit(ctx: Context, old_msg: Message, new_msg: Message) -> bool: +def predicate_message_edit(ctx: Context, old_msg: Message, new_msg: Message) -> bool: """Return True if the edited message is the context message and the content was indeed modified.""" return new_msg.id == ctx.message.id and old_msg.content != new_msg.content -def predicate_eval_emoji_reaction(ctx: Context, reaction: Reaction, user: User) -> bool: - """Return True if the reaction REEVAL_EMOJI was added by the context message author on this message.""" - return reaction.message.id == ctx.message.id and user.id == ctx.author.id and str(reaction) == REEVAL_EMOJI +def predicate_emoji_reaction(ctx: Context, reaction: Reaction, user: User) -> bool: + """Return True if the reaction REDO_EMOJI was added by the context message author on this message.""" + return reaction.message.id == ctx.message.id and user.id == ctx.author.id and str(reaction) == REDO_EMOJI def setup(bot: Bot) -> None: diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index 339cdaaa4..5d213a883 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -209,7 +209,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.format_output = AsyncMock(return_value=('[No output]', None)) mocked_filter_cog = MagicMock() - mocked_filter_cog.filter_eval = AsyncMock(return_value=False) + mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') @@ -241,7 +241,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.format_output = AsyncMock(return_value=('Way too long beard', 'lookatmybeard.com')) mocked_filter_cog = MagicMock() - mocked_filter_cog.filter_eval = AsyncMock(return_value=False) + mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') @@ -270,7 +270,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.format_output = AsyncMock() # This function isn't called mocked_filter_cog = MagicMock() - mocked_filter_cog.filter_eval = AsyncMock(return_value=False) + mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') -- cgit v1.2.3 From 85a6f430aa68f59ce6958ecb6450eca0736628e4 Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Thu, 27 Jan 2022 10:31:12 -0500 Subject: chore: Switch Snekbox.prepare_input with a CodeblockConverter As per @Numerlor's suggestion --- bot/exts/filters/filtering.py | 2 +- bot/exts/utils/snekbox.py | 74 ++++++++++++------------ tests/bot/exts/utils/test_snekbox.py | 105 +++++++++++++++++------------------ 3 files changed, 91 insertions(+), 90 deletions(-) (limited to 'tests') diff --git a/bot/exts/filters/filtering.py b/bot/exts/filters/filtering.py index 599302576..375e9dca8 100644 --- a/bot/exts/filters/filtering.py +++ b/bot/exts/filters/filtering.py @@ -267,7 +267,7 @@ class Filtering(Cog): # Update time when alert sent await self.name_alerts.set(member.id, arrow.utcnow().timestamp()) - async def filter_snekbox_job(self, result: str, msg: Message) -> bool: + async def filter_snekbox_output(self, result: str, msg: Message) -> bool: """ Filter the result of a snekbox command to see if it violates any of our rules, and then respond accordingly. diff --git a/bot/exts/utils/snekbox.py b/bot/exts/utils/snekbox.py index 15599208f..41f6bf8ad 100644 --- a/bot/exts/utils/snekbox.py +++ b/bot/exts/utils/snekbox.py @@ -9,7 +9,7 @@ from typing import Optional, Tuple from botcore.regex import FORMATTED_CODE_REGEX, RAW_CODE_REGEX from discord import AllowedMentions, HTTPException, Message, NotFound, Reaction, User -from discord.ext.commands import Cog, Command, Context, command, guild_only +from discord.ext.commands import Cog, Command, Context, Converter, command, guild_only from bot.bot import Bot from bot.constants import Categories, Channels, Roles, URLs @@ -68,35 +68,11 @@ REDO_EMOJI = '\U0001f501' # :repeat: REDO_TIMEOUT = 30 -class Snekbox(Cog): - """Safe evaluation of Python code using Snekbox.""" - - def __init__(self, bot: Bot): - self.bot = bot - self.jobs = {} - - async def post_job(self, code: str, *, args: Optional[list[str]] = None) -> dict: - """Send a POST request to the Snekbox API to evaluate code and return the results.""" - url = URLs.snekbox_eval_api - data = {"input": code} - - if args is not None: - data["args"] = args +class CodeblockConverter(Converter): + """Attempts to extract code from a codeblock, if provided.""" - async with self.bot.http_session.post(url, json=data, raise_for_status=True) as resp: - return await resp.json() - - async def upload_output(self, output: str) -> Optional[str]: - """Upload the job's output to a paste service and return a URL to it if successful.""" - log.trace("Uploading full output to paste service...") - - if len(output) > MAX_PASTE_LEN: - log.info("Full output is too long to upload") - return "too long to upload" - return await send_to_paste_service(output, extension="txt") - - @staticmethod - def prepare_input(code: str) -> list[str]: + @classmethod + async def convert(cls, ctx: Context, code: str) -> list[str]: """ Extract code from the Markdown, format it, and insert it into the code template. @@ -128,6 +104,34 @@ class Snekbox(Cog): log.trace(f"Extracted {info} for evaluation:\n{code}") return codeblocks + +class Snekbox(Cog): + """Safe evaluation of Python code using Snekbox.""" + + def __init__(self, bot: Bot): + self.bot = bot + self.jobs = {} + + async def post_job(self, code: str, *, args: Optional[list[str]] = None) -> dict: + """Send a POST request to the Snekbox API to evaluate code and return the results.""" + url = URLs.snekbox_eval_api + data = {"input": code} + + if args is not None: + data["args"] = args + + async with self.bot.http_session.post(url, json=data, raise_for_status=True) as resp: + return await resp.json() + + async def upload_output(self, output: str) -> Optional[str]: + """Upload the job's output to a paste service and return a URL to it if successful.""" + log.trace("Uploading full output to paste service...") + + if len(output) > MAX_PASTE_LEN: + log.info("Full output is too long to upload") + return "too long to upload" + return await send_to_paste_service(output, extension="txt") + @staticmethod def prepare_timeit_input(codeblocks: list[str]) -> tuple[str, list[str]]: """ @@ -313,7 +317,7 @@ class Snekbox(Cog): await ctx.message.clear_reaction(REDO_EMOJI) return None, None - codeblocks = self.prepare_input(code) + codeblocks = await CodeblockConverter.convert(ctx, code) if command is self.timeit_command: return self.prepare_timeit_input(codeblocks) @@ -393,7 +397,7 @@ class Snekbox(Cog): channels=NO_SNEKBOX_CHANNELS, ping_user=False ) - async def eval_command(self, ctx: Context, *, code: str) -> None: + async def eval_command(self, ctx: Context, *, code: CodeblockConverter) -> None: """ Run Python code and get the results. @@ -404,8 +408,7 @@ class Snekbox(Cog): We've done our best to make this sandboxed, but do let us know if you manage to find an issue with it! """ - code = "\n".join(self.prepare_input(code)) - await self.run_job("eval", ctx, code) + await self.run_job("eval", ctx, "\n".join(code)) @command(name="timeit", aliases=("ti",)) @guild_only() @@ -416,7 +419,7 @@ class Snekbox(Cog): channels=NO_SNEKBOX_CHANNELS, ping_user=False ) - async def timeit_command(self, ctx: Context, *, code: str) -> str: + async def timeit_command(self, ctx: Context, *, code: CodeblockConverter) -> str: """ Profile Python Code to find execution time. @@ -430,8 +433,7 @@ class Snekbox(Cog): We've done our best to make this sandboxed, but do let us know if you manage to find an issue with it! """ - codeblocks = self.prepare_input(code) - code, args = self.prepare_timeit_input(codeblocks) + code, args = self.prepare_timeit_input(code) await self.run_job("timeit", ctx, code=code, args=args) diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index 5d213a883..75da0c860 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -17,7 +17,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.bot = MockBot() self.cog = Snekbox(bot=self.bot) - async def test_post_eval(self): + async def test_post_job(self): """Post the eval code to the URLs.snekbox_eval_api endpoint.""" resp = MagicMock() resp.json = AsyncMock(return_value="return") @@ -26,7 +26,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): context_manager.__aenter__.return_value = resp self.bot.http_session.post.return_value = context_manager - self.assertEqual(await self.cog.post_eval("import random"), "return") + self.assertEqual(await self.cog.post_job("import random"), "return") self.bot.http_session.post.assert_called_with( constants.URLs.snekbox_eval_api, json={"input": "import random"}, @@ -45,7 +45,8 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): await self.cog.upload_output("Test output.") mock_paste_util.assert_called_once_with("Test output.", extension="txt") - def test_prepare_input(self): + async def test_codeblock_converter(self): + ctx = MockContext() cases = ( ('print("Hello world!")', 'print("Hello world!")', 'non-formatted'), ('`print("Hello world!")`', 'print("Hello world!")', 'one line code block'), @@ -61,7 +62,9 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): ) for case, expected, testname in cases: with self.subTest(msg=f'Extract code from {testname}.'): - self.assertEqual('\n'.join(self.cog.prepare_input(case)), expected) + self.assertEqual( + '\n'.join(await snekbox.CodeblockConverter.convert(ctx, case)), expected + ) def test_get_results_message(self): """Return error and message according to the eval result.""" @@ -158,31 +161,27 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): response = MockMessage() ctx.command = MagicMock() - self.cog.prepare_input = MagicMock(return_value=['MyAwesomeFormattedCode']) - self.cog.send_eval = AsyncMock(return_value=response) - self.cog.continue_eval = AsyncMock(return_value=(None, None)) + self.cog.send_job = AsyncMock(return_value=response) + self.cog.continue_job = AsyncMock(return_value=(None, None)) - await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') - self.cog.prepare_input.assert_called_once_with('MyAwesomeCode') - self.cog.send_eval.assert_called_once_with(ctx, 'MyAwesomeFormattedCode', args=None, job_name='eval') - self.cog.continue_eval.assert_called_once_with(ctx, response, ctx.command) + await self.cog.eval_command(self.cog, ctx=ctx, code=['MyAwesomeCode']) + self.cog.send_job.assert_called_once_with(ctx, 'MyAwesomeCode', args=None, job_name='eval') + self.cog.continue_job.assert_called_once_with(ctx, response, ctx.command) async def test_eval_command_evaluate_twice(self): """Test the eval and re-eval command procedure.""" ctx = MockContext() response = MockMessage() ctx.command = MagicMock() - self.cog.prepare_input = MagicMock(return_value='MyAwesomeFormattedCode') - self.cog.send_eval = AsyncMock(return_value=response) - self.cog.continue_eval = AsyncMock() - self.cog.continue_eval.side_effect = (('MyAwesomeFormattedCode', None), (None, None)) + self.cog.send_job = AsyncMock(return_value=response) + self.cog.continue_job = AsyncMock() + self.cog.continue_job.side_effect = (('MyAwesomeFormattedCode', None), (None, None)) - await self.cog.eval_command(self.cog, ctx=ctx, code='MyAwesomeCode') - self.cog.prepare_input.has_calls(call('MyAwesomeCode'), call('MyAwesomeCode-2')) - self.cog.send_eval.assert_called_with( + await self.cog.eval_command(self.cog, ctx=ctx, code=['MyAwesomeCode']) + self.cog.send_job.assert_called_with( ctx, 'MyAwesomeFormattedCode', args=None, job_name='eval' ) - self.cog.continue_eval.assert_called_with(ctx, response, ctx.command) + self.cog.continue_job.assert_called_with(ctx, response, ctx.command) async def test_eval_command_reject_two_eval_at_the_same_time(self): """Test if the eval command rejects an eval if the author already have a running eval.""" @@ -196,14 +195,14 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): "@LemonLemonishBeard#0042 You've already got a job running - please wait for it to finish!" ) - async def test_send_eval(self): - """Test the send_eval function.""" + async def test_send_job(self): + """Test the send_job function.""" ctx = MockContext() ctx.message = MockMessage() ctx.send = AsyncMock() ctx.author = MockUser(mention='@LemonLemonishBeard#0042') - self.cog.post_eval = AsyncMock(return_value={'stdout': '', 'returncode': 0}) + self.cog.post_job = AsyncMock(return_value={'stdout': '', 'returncode': 0}) self.cog.get_results_message = MagicMock(return_value=('Return code 0', '')) self.cog.get_status_emoji = MagicMock(return_value=':yay!:') self.cog.format_output = AsyncMock(return_value=('[No output]', None)) @@ -212,7 +211,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') + await self.cog.send_job(ctx, 'MyAwesomeCode', job_name='eval') ctx.send.assert_called_once() self.assertEqual( @@ -223,19 +222,19 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): expected_allowed_mentions = AllowedMentions(everyone=False, roles=False, users=[ctx.author]) self.assertEqual(allowed_mentions.to_dict(), expected_allowed_mentions.to_dict()) - self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) + self.cog.post_job.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': '', 'returncode': 0}) self.cog.get_results_message.assert_called_once_with({'stdout': '', 'returncode': 0}, 'eval') self.cog.format_output.assert_called_once_with('') - async def test_send_eval_with_paste_link(self): - """Test the send_eval function with a too long output that generate a paste link.""" + async def test_send_job_with_paste_link(self): + """Test the send_job function with a too long output that generate a paste link.""" ctx = MockContext() ctx.message = MockMessage() ctx.send = AsyncMock() ctx.author.mention = '@LemonLemonishBeard#0042' - self.cog.post_eval = AsyncMock(return_value={'stdout': 'Way too long beard', 'returncode': 0}) + self.cog.post_job = AsyncMock(return_value={'stdout': 'Way too long beard', 'returncode': 0}) self.cog.get_results_message = MagicMock(return_value=('Return code 0', '')) self.cog.get_status_emoji = MagicMock(return_value=':yay!:') self.cog.format_output = AsyncMock(return_value=('Way too long beard', 'lookatmybeard.com')) @@ -244,7 +243,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') + await self.cog.send_job(ctx, 'MyAwesomeCode', job_name='eval') ctx.send.assert_called_once() self.assertEqual( @@ -253,18 +252,18 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): '\n\n```\nWay too long beard\n```\nFull output: lookatmybeard.com' ) - self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) + self.cog.post_job.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': 'Way too long beard', 'returncode': 0}) self.cog.get_results_message.assert_called_once_with({'stdout': 'Way too long beard', 'returncode': 0}, 'eval') self.cog.format_output.assert_called_once_with('Way too long beard') - async def test_send_eval_with_non_zero_eval(self): - """Test the send_eval function with a code returning a non-zero code.""" + async def test_send_job_with_non_zero_eval(self): + """Test the send_job function with a code returning a non-zero code.""" ctx = MockContext() ctx.message = MockMessage() ctx.send = AsyncMock() ctx.author.mention = '@LemonLemonishBeard#0042' - self.cog.post_eval = AsyncMock(return_value={'stdout': 'ERROR', 'returncode': 127}) + self.cog.post_job = AsyncMock(return_value={'stdout': 'ERROR', 'returncode': 127}) self.cog.get_results_message = MagicMock(return_value=('Return code 127', 'Beard got stuck in the eval')) self.cog.get_status_emoji = MagicMock(return_value=':nope!:') self.cog.format_output = AsyncMock() # This function isn't called @@ -273,7 +272,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog - await self.cog.send_eval(ctx, 'MyAwesomeCode', job_name='eval') + await self.cog.send_job(ctx, 'MyAwesomeCode', job_name='eval') ctx.send.assert_called_once() self.assertEqual( @@ -281,14 +280,14 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): '@LemonLemonishBeard#0042 :nope!: Return code 127.\n\n```\nBeard got stuck in the eval\n```' ) - self.cog.post_eval.assert_called_once_with('MyAwesomeCode', args=None) + self.cog.post_job.assert_called_once_with('MyAwesomeCode', args=None) self.cog.get_status_emoji.assert_called_once_with({'stdout': 'ERROR', 'returncode': 127}) self.cog.get_results_message.assert_called_once_with({'stdout': 'ERROR', 'returncode': 127}, 'eval') self.cog.format_output.assert_not_called() @patch("bot.exts.utils.snekbox.partial") - async def test_continue_eval_does_continue(self, partial_mock): - """Test that the continue_eval function does continue if required conditions are met.""" + async def test_continue_job_does_continue(self, partial_mock): + """Test that the continue_job function does continue if required conditions are met.""" ctx = MockContext(message=MockMessage(add_reaction=AsyncMock(), clear_reactions=AsyncMock())) response = MockMessage(delete=AsyncMock()) new_msg = MockMessage() @@ -296,30 +295,30 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): expected = "NewCode" self.cog.get_code = create_autospec(self.cog.get_code, spec_set=True, return_value=expected) - actual = await self.cog.continue_eval(ctx, response, self.cog.eval_command) + actual = await self.cog.continue_job(ctx, response, self.cog.eval_command) self.cog.get_code.assert_awaited_once_with(new_msg, ctx.command) self.assertEqual(actual, (expected, None)) self.bot.wait_for.assert_has_awaits( ( call( 'message_edit', - check=partial_mock(snekbox.predicate_eval_message_edit, ctx), - timeout=snekbox.REEVAL_TIMEOUT, + check=partial_mock(snekbox.predicate_message_edit, ctx), + timeout=snekbox.REDO_TIMEOUT, ), - call('reaction_add', check=partial_mock(snekbox.predicate_eval_emoji_reaction, ctx), timeout=10) + call('reaction_add', check=partial_mock(snekbox.predicate_emoji_reaction, ctx), timeout=10) ) ) - ctx.message.add_reaction.assert_called_once_with(snekbox.REEVAL_EMOJI) - ctx.message.clear_reaction.assert_called_once_with(snekbox.REEVAL_EMOJI) + ctx.message.add_reaction.assert_called_once_with(snekbox.REDO_EMOJI) + ctx.message.clear_reaction.assert_called_once_with(snekbox.REDO_EMOJI) response.delete.assert_called_once() - async def test_continue_eval_does_not_continue(self): + async def test_continue_job_does_not_continue(self): ctx = MockContext(message=MockMessage(clear_reactions=AsyncMock())) self.bot.wait_for.side_effect = asyncio.TimeoutError - actual = await self.cog.continue_eval(ctx, MockMessage(), self.cog.eval_command) + actual = await self.cog.continue_job(ctx, MockMessage(), self.cog.eval_command) self.assertEqual(actual, (None, None)) - ctx.message.clear_reaction.assert_called_once_with(snekbox.REEVAL_EMOJI) + ctx.message.clear_reaction.assert_called_once_with(snekbox.REDO_EMOJI) async def test_get_code(self): """Should return 1st arg (or None) if eval cmd in message, otherwise return full content.""" @@ -347,8 +346,8 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.bot.get_context.assert_awaited_once_with(message) self.assertEqual(actual_code, expected_code) - def test_predicate_eval_message_edit(self): - """Test the predicate_eval_message_edit function.""" + def test_predicate_message_edit(self): + """Test the predicate_message_edit function.""" msg0 = MockMessage(id=1, content='abc') msg1 = MockMessage(id=2, content='abcdef') msg2 = MockMessage(id=1, content='abcdef') @@ -361,18 +360,18 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): for ctx_msg, new_msg, expected, testname in cases: with self.subTest(msg=f'Messages with {testname} return {expected}'): ctx = MockContext(message=ctx_msg) - actual = snekbox.predicate_eval_message_edit(ctx, ctx_msg, new_msg) + actual = snekbox.predicate_message_edit(ctx, ctx_msg, new_msg) self.assertEqual(actual, expected) - def test_predicate_eval_emoji_reaction(self): - """Test the predicate_eval_emoji_reaction function.""" + def test_predicate_emoji_reaction(self): + """Test the predicate_emoji_reaction function.""" valid_reaction = MockReaction(message=MockMessage(id=1)) - valid_reaction.__str__.return_value = snekbox.REEVAL_EMOJI + valid_reaction.__str__.return_value = snekbox.REDO_EMOJI valid_ctx = MockContext(message=MockMessage(id=1), author=MockUser(id=2)) valid_user = MockUser(id=2) invalid_reaction_id = MockReaction(message=MockMessage(id=42)) - invalid_reaction_id.__str__.return_value = snekbox.REEVAL_EMOJI + invalid_reaction_id.__str__.return_value = snekbox.REDO_EMOJI invalid_user_id = MockUser(id=42) invalid_reaction_str = MockReaction(message=MockMessage(id=1)) invalid_reaction_str.__str__.return_value = ':longbeard:' @@ -385,7 +384,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): ) for reaction, user, expected, testname in cases: with self.subTest(msg=f'Test with {testname} and expected return {expected}'): - actual = snekbox.predicate_eval_emoji_reaction(valid_ctx, reaction, user) + actual = snekbox.predicate_emoji_reaction(valid_ctx, reaction, user) self.assertEqual(actual, expected) -- cgit v1.2.3 From b689f059e45e68f75e3a97277ac3bf58263fa769 Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Fri, 4 Feb 2022 20:54:49 -0500 Subject: fix: Use filter_snekbox_output rather than job --- bot/exts/utils/snekbox.py | 2 +- tests/bot/exts/utils/test_snekbox.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) (limited to 'tests') diff --git a/bot/exts/utils/snekbox.py b/bot/exts/utils/snekbox.py index ce3dd7c24..a932b96ff 100644 --- a/bot/exts/utils/snekbox.py +++ b/bot/exts/utils/snekbox.py @@ -268,7 +268,7 @@ class Snekbox(Cog): filter_cog = self.bot.get_cog("Filtering") filter_triggered = False if filter_cog: - filter_triggered = await filter_cog.filter_snekbox_job(msg, ctx.message) + filter_triggered = await filter_cog.filter_snekbox_output(msg, ctx.message) if filter_triggered: response = await ctx.send("Attempt to circumvent filter detected. Moderator team has been alerted.") else: diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index 75da0c860..2eaed0446 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -208,7 +208,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.format_output = AsyncMock(return_value=('[No output]', None)) mocked_filter_cog = MagicMock() - mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) + mocked_filter_cog.filter_snekbox_output = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog await self.cog.send_job(ctx, 'MyAwesomeCode', job_name='eval') @@ -240,7 +240,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.format_output = AsyncMock(return_value=('Way too long beard', 'lookatmybeard.com')) mocked_filter_cog = MagicMock() - mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) + mocked_filter_cog.filter_snekbox_output = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog await self.cog.send_job(ctx, 'MyAwesomeCode', job_name='eval') @@ -269,7 +269,7 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): self.cog.format_output = AsyncMock() # This function isn't called mocked_filter_cog = MagicMock() - mocked_filter_cog.filter_snekbox_job = AsyncMock(return_value=False) + mocked_filter_cog.filter_snekbox_output = AsyncMock(return_value=False) self.bot.get_cog.return_value = mocked_filter_cog await self.cog.send_job(ctx, 'MyAwesomeCode', job_name='eval') -- cgit v1.2.3 From 7ef9b2163e7f0da5a7abd17decac2b0b2d53defb Mon Sep 17 00:00:00 2001 From: ToxicKidz Date: Sat, 5 Feb 2022 23:21:21 -0500 Subject: tests: Add a test for timeit command codeblock preparation --- tests/bot/exts/utils/test_snekbox.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) (limited to 'tests') diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py index 2eaed0446..f68a20089 100644 --- a/tests/bot/exts/utils/test_snekbox.py +++ b/tests/bot/exts/utils/test_snekbox.py @@ -66,6 +66,21 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase): '\n'.join(await snekbox.CodeblockConverter.convert(ctx, case)), expected ) + def test_prepare_timeit_input(self): + """Test the prepare_timeit_input codeblock detection.""" + base_args = ('-m', 'timeit', '-s') + cases = ( + (['print("Hello World")'], '', 'single block of code'), + (['x = 1', 'print(x)'], 'x = 1', 'two blocks of code'), + (['x = 1', 'print(x)', 'print("Some other code.")'], 'x = 1', 'three blocks of code') + ) + + for case, setup_code, testname in cases: + setup = snekbox.TIMEIT_SETUP_WRAPPER.format(setup=setup_code) + expected = ('\n'.join(case[1:] if setup_code else case), [*base_args, setup]) + with self.subTest(msg=f'Test with {testname} and expected return {expected}'): + self.assertEqual(self.cog.prepare_timeit_input(case), expected) + def test_get_results_message(self): """Return error and message according to the eval result.""" cases = ( -- cgit v1.2.3