Add unit tests for plot capture, upload, and error handling in evaluation function

Introduces tests for single and multiple plots, graceful upload failure handling, and `_run_code` image capture. Adjusts `_run_code` to ensure plot finalization via `_capture_plots` call.

Author: m-messer

evaluation_function/evaluation.py

@@ -13,7 +13,6 @@
 
 _PREAMBLE_TEMPLATE = """\
 import os as _os
-import atexit as _atexit
 
 _plot_dir = {plot_dir!r}
 _plot_idx = [0]
@@ -27,16 +26,16 @@ def _capture_plots():
         _plot_idx[0] += 1
         _plt.figure(num).savefig(
             _os.path.join(_plot_dir, str(_plot_idx[0]).zfill(4) + '.png'))
-
-_atexit.register(_capture_plots)
 """
 
+_CAPTURE_CALL = "_capture_plots()\n"
+
 
 def _run_code(code: str, stdin: str) -> tuple[str, str, bool, list[Image.Image]]:
     plot_dir = tempfile.mkdtemp()
     preamble = _PREAMBLE_TEMPLATE.format(plot_dir=plot_dir)
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
-        f.write(preamble + "\n" + code)
+        f.write(preamble + "\n" + code + "\n" + _CAPTURE_CALL)
         tmpfile = f.name
     try:
         proc = subprocess.run(

evaluation_function/evaluation_test.py

@@ -1,6 +1,9 @@
 import unittest
+from unittest.mock import patch
 
-from .evaluation import evaluation_function
+from lf_toolkit.evaluation.image_upload import ImageUploadError
+
+from .evaluation import evaluation_function, _run_code
 
 _SQUARE_CODE = "n = int(input())\nprint(n * n)"
 _CRASH_CODE = "raise ValueError('oops')"
@@ -54,4 +57,52 @@ def test_no_tests(self):
         result = evaluation_function(_SQUARE_CODE, None, {}).to_dict()
 
         self.assertFalse(result["is_correct"])
-        self.assertIn("```", result["feedback"])
+        self.assertIn("```", result["feedback"])
+
+
+_PLOT_CODE = "import matplotlib.pyplot as plt\nplt.plot([1, 2, 3])\n"
+_MULTI_PLOT_CODE = (
+    "import matplotlib.pyplot as plt\n"
+    "plt.figure(1); plt.plot([1, 2])\n"
+    "plt.figure(2); plt.plot([3, 4])\n"
+)
+_FAKE_URL = "https://example.com/plot.png"
+
+
+class TestImageGeneration(unittest.TestCase):
+
+    @patch("evaluation_function.evaluation.upload_image", return_value=_FAKE_URL)
+    def test_single_plot_demo_mode(self, mock_upload):
+        result = evaluation_function(_PLOT_CODE, None, {}).to_dict()
+        mock_upload.assert_called_once()
+        self.assertIn("![Plot 1]", result["feedback"])
+
+    @patch("evaluation_function.evaluation.upload_image", return_value=_FAKE_URL)
+    def test_single_plot_passing_test(self, mock_upload):
+        params = _params(_test("", ""))
+        result = evaluation_function(_PLOT_CODE, None, params).to_dict()
+        self.assertIn("![Plot 1]", result["feedback"])
+
+    @patch("evaluation_function.evaluation.upload_image", return_value=_FAKE_URL)
+    def test_single_plot_failing_test(self, mock_upload):
+        params = _params(_test("", "wrong"))
+        result = evaluation_function(_PLOT_CODE, None, params).to_dict()
+        self.assertIn("![Plot 1]", result["feedback"])
+
+    @patch("evaluation_function.evaluation.upload_image", return_value=_FAKE_URL)
+    def test_multiple_plots(self, mock_upload):
+        result = evaluation_function(_MULTI_PLOT_CODE, None, {}).to_dict()
+        self.assertEqual(mock_upload.call_count, 2)
+        self.assertIn("![Plot 1]", result["feedback"])
+        self.assertIn("![Plot 2]", result["feedback"])
+
+    @patch("evaluation_function.evaluation.upload_image", side_effect=ImageUploadError)
+    def test_upload_failure_graceful(self, mock_upload):
+        result = evaluation_function(_PLOT_CODE, None, {}).to_dict()
+        self.assertNotIn("![Plot", result["feedback"])
+
+    def test_run_code_captures_images(self):
+        _, _, timed_out, images = _run_code(_PLOT_CODE, "")
+        self.assertFalse(timed_out)
+        self.assertEqual(len(images), 1)
+        self.assertEqual(images[0].format, "PNG")