diff --git a/tutorial/tests/testsuite/helpers.py b/tutorial/tests/testsuite/helpers.py
index 5c7f17c5..4dddf1e1 100644
--- a/tutorial/tests/testsuite/helpers.py
+++ b/tutorial/tests/testsuite/helpers.py
@@ -11,6 +11,7 @@
 from IPython.display import Code
 from IPython.display import display as ipython_display
 from ipywidgets import HTML
+from testsuite import iPythonGenerator
 
 from .ai_helpers import AIExplanation, OpenAIWrapper
 
@@ -377,6 +378,9 @@ class IPytestResult:
 class TestResultOutput:
     """Class to prepare and display test results in a Jupyter notebook"""
 
+    def __init__(self, output_generator=iPythonGenerator()) -> None:
+        self.output_generator = output_generator
+
     ipytest_result: IPytestResult
     solution: Optional[str] = None
     MAX_ATTEMPTS: ClassVar[int] = 3
@@ -409,29 +413,12 @@ def display_results(self) -> None:
                     self.MAX_ATTEMPTS - self.ipytest_result.test_attempts
                 )
                 cells.append(
-                    HTML(
-                        '<div style="margin-top: 1.5rem; font-family: system-ui, -apple-system, sans-serif;">'
-                        f'<div style="display: flex; align-items: center; gap: 0.5rem;">'
-                        '<span style="font-size: 1.2rem;">📝</span>'
-                        '<span style="font-size: 1.1rem; font-weight: 500;">Solution will be available after '
-                        f'{attempts_remaining} more failed attempt{"s" if attempts_remaining > 1 else ""}</span>'
-                        "</div>"
-                        "</div>"
+                    self.output_generator.attempts_remaining_explanation(
+                        attempts_remaining
                     )
                 )
 
-        ipython_display(
-            ipywidgets.VBox(
-                children=cells,
-                layout={
-                    "border": "1px solid #e5e7eb",
-                    "background-color": "#ffffff",
-                    "margin": "5px",
-                    "padding": "0.75rem",
-                    "border-radius": "0.5rem",
-                },
-            )
-        )
+        iPythonGenerator.display_cells(cells)
 
     # TODO: This is left for reference if we ever want to bring back this styling
     # Perhaps we should remove it if it's unnecessary
diff --git a/tutorial/tests/testsuite/testsuite.py b/tutorial/tests/testsuite/testsuite.py
index 2bb8398f..b18b5240 100644
--- a/tutorial/tests/testsuite/testsuite.py
+++ b/tutorial/tests/testsuite/testsuite.py
@@ -6,6 +6,7 @@
 import io
 import os
 import pathlib
+from abc import ABC, abstractmethod
 from collections import defaultdict
 from contextlib import contextmanager, redirect_stderr, redirect_stdout
 from queue import Queue
@@ -13,11 +14,13 @@
 from typing import Dict, List, Optional
 
 import ipynbname
+import ipywidgets
 import pytest
 from dotenv import find_dotenv, load_dotenv
 from IPython.core.interactiveshell import InteractiveShell
 from IPython.core.magic import Magics, cell_magic, magics_class
-from IPython.display import HTML, display
+from IPython.display import HTML
+from IPython.display import display as ipython_display
 
 from .ai_helpers import OpenAIWrapper
 from .ast_parser import AstParser
@@ -143,11 +146,58 @@ def get_module_name(line: str, globals_dict: Dict) -> str | None:
     return module_name
 
 
+class OutputGenerator(ABC):
+    @abstractmethod
+    def display(self, content: str) -> None:
+        pass
+
+    @abstractmethod
+    def attempts_remaining_explanation(attempts_remaining: int) -> str:
+        pass
+
+
+class iPythonGenerator(OutputGenerator):
+    def attempts_remaining_explanation(attempts_remaining: int) -> str:
+        return (
+            '<div style="margin-top: 1.5rem; font-family: system-ui, -apple-system, sans-serif;">'
+            f'<div style="display: flex; align-items: center; gap: 0.5rem;">'
+            '<span style="font-size: 1.2rem;">📝</span>'
+            '<span style="font-size: 1.1rem; font-weight: 500;">Solution will be available after '
+            f'{attempts_remaining} more failed attempt{"s" if attempts_remaining > 1 else ""}</span>'
+            "</div>"
+            "</div>"
+        )
+
+    def display_cells(cells):
+        html_cells = [HTML(cell) for cell in cells]
+        ipython_display(
+            ipywidgets.VBox(
+                children=html_cells,
+                layout={
+                    "border": "1px solid #e5e7eb",
+                    "background-color": "#ffffff",
+                    "margin": "5px",
+                    "padding": "0.75rem",
+                    "border-radius": "0.5rem",
+                },
+            )
+        )
+        return None
+
+
+class DebugGenerator(OutputGenerator):
+    def attempts_remaining_explanation(attempts_remaining: int) -> str:
+        return f"{attempts_remaining}"
+
+    def display_cells(cells):
+        return cells
+
+
 @magics_class
 class TestMagic(Magics):
     """Class to add the test cell magic"""
 
-    def __init__(self, shell):
+    def __init__(self, shell, output_generator=iPythonGenerator()):
         super().__init__(shell)
         self.shell: InteractiveShell = shell
         self.cell: str = ""
@@ -160,6 +210,7 @@ def __init__(self, shell):
         )
         self._orig_traceback = self.shell._showtraceback  # type: ignore
         # This is monkey-patching suppress printing any exception or traceback
+        self.output_generator = output_generator
 
     def extract_functions_to_test(self) -> List[AFunction]:
         """Retrieve the functions names and implementations defined in the current cell"""
@@ -311,7 +362,7 @@ def ipytest(self, line: str, cell: str):
                     module_file=self.module_file,
                     results=results,
                 )
-                display(HTML(debug_output.to_html()))
+                output = self.output_generator.display(debug_output)
 
             # Parse the AST of the test module to retrieve the solution code
             ast_parser = AstParser(self.module_file)
@@ -329,7 +380,7 @@ def ipytest(self, line: str, cell: str):
                 ).display_results()
 
 
-def load_ipython_extension(ipython):
+def load_ipython_extension(ipython, output_generator):
     """
     Any module file that define a function named `load_ipython_extension`
     can be loaded via `%load_ext module.path` or be configured to be
@@ -377,13 +428,11 @@ def load_ipython_extension(ipython):
             )
             message_color = "#ffebee"
 
-    display(
-        HTML(
-            "<div style='background-color: "
-            f"{message_color}; border-radius: 5px; padding: 10px;'>"
-            f"{message}"
-            "</div>"
-        )
+    output_generator.display(
+        "<div style='background-color: "
+        f"{message_color}; border-radius: 5px; padding: 10px;'>"
+        f"{message}"
+        "</div>"
     )
 
     # Register the magic
@@ -393,4 +442,4 @@ def load_ipython_extension(ipython):
         "<div style='background-color: #fffde7; border-radius: 5px; padding: 10px;'>"
         "🔄 <strong>IPytest extension (re)loaded.</strong></div>"
     )
-    display(HTML(message))
+    output_generator.display(message)