Standardisation 8

tests/test_gui_core.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import importlib.util
 import io
+import json
 import os
 import queue
 import sys
@@ -11,6 +12,7 @@ import unittest
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import redirect_stderr
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 from PIL import Image
@@ -21,6 +23,7 @@ from copienator import (
     WorkspaceNotFoundError,
     WorkspaceValidationError,
     atomic_update_json,
+    atomic_write_bytes,
     atomic_write_json,
     read_json,
     workspace_from_target,
@@ -114,6 +117,12 @@ class WorkspaceTests(unittest.TestCase):
 
 
 class AtomicJsonTests(unittest.TestCase):
+    def test_atomic_binary_round_trip(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "result.jsonl"
+            atomic_write_bytes(path, b'{"one":1}\n')
+            self.assertEqual(path.read_bytes(), b'{"one":1}\n')
+
     def test_atomic_round_trip_and_unicode(self) -> None:
         with tempfile.TemporaryDirectory() as directory:
             path = Path(directory) / "nested" / "state.json"
@@ -313,6 +322,14 @@ class StandardCliTests(unittest.TestCase):
             "gemini_for_labels": load_script_module(
                 "gemini_for_labels.py", "gemini_for_labels"
             ),
+            "correction": load_script_module("correction.py", "correction"),
+            "submit_batches": load_script_module(
+                "submit_batches.py", "submit_batches"
+            ),
+            "batch_status": load_script_module("batch_status.py", "batch_status"),
+            "fetch_batched_results": load_script_module(
+                "fetch_batched_results.py", "fetch_batched_results"
+            ),
             "copies_tools": load_script_module(
                 "copies_tools.py", "copienator_copies_tools_test"
             ),
@@ -355,6 +372,9 @@ class StandardCliTests(unittest.TestCase):
                 "page_splitter": [missing],
                 "plotting": [missing],
                 "gemini_for_labels": [missing],
+                "correction": [missing],
+                "submit_batches": [missing],
+                "fetch_batched_results": [missing],
             }
             for name, arguments in invocations.items():
                 with self.subTest(script=name), redirect_stderr(io.StringIO()):
@@ -472,6 +492,21 @@ class StandardCliTests(unittest.TestCase):
                 "default",
                 {"target": evaluation, "overwrite": True},
             ),
+            "correction": (
+                "correction",
+                "live",
+                {"target": evaluation, "overwrite": True, "limit": 5},
+            ),
+            "submit_batches": (
+                "submit_batches",
+                "default",
+                {"target": evaluation},
+            ),
+            "fetch_batched_results": (
+                "fetch_batches",
+                "default",
+                {"target": evaluation},
+            ),
         }
         for module_name, (step_id, variant_id, values) in cases.items():
             step = steps[step_id]
@@ -900,6 +935,135 @@ class StandardCliTests(unittest.TestCase):
                 "Ex 1",
             )
 
+    def test_correction_overwrite_keeps_previous_state_until_a_commit(self) -> None:
+        module = self.modules["correction"]
+        with tempfile.TemporaryDirectory() as directory:
+            evaluation = Path(directory) / "Exam"
+            evaluation.mkdir()
+            correction = evaluation / "correction.json"
+            progress = evaluation / "correction_progress.json"
+            atomic_write_json(correction, {"Ex 1": [[{"id": "01"}]]})
+            atomic_write_json(progress, [["old.jpg", "Ex 1"]])
+            args = module.build_parser().parse_args(
+                [str(evaluation), "--overwrite"]
+            )
+            module.configure_runtime(
+                EvaluationWorkspace(evaluation),
+                [("new.jpg", "Ex 1")],
+                args,
+                api_client=Mock(),
+            )
+            self.assertEqual(read_json(correction), {"Ex 1": [[{"id": "01"}]]})
+            self.assertEqual(read_json(progress), [["old.jpg", "Ex 1"]])
+            self.assertEqual(module.results, {"Ex 1": []})
+            self.assertEqual(module.tasks_to_process, [("new.jpg", "Ex 1")])
+
+    def test_correction_reset_restores_old_and_deletes_new_files(self) -> None:
+        module = self.modules["correction"]
+        with tempfile.TemporaryDirectory() as directory:
+            evaluation = Path(directory) / "Exam"
+            copy_dir = evaluation / "Copies" / "Copie01"
+            copy_dir.mkdir(parents=True)
+            atomic_write_json(evaluation / "correction.json", {"old": True})
+            atomic_write_json(evaluation / "correction_progress.json", ["old"])
+            (copy_dir / "Ex 1.pdf").write_bytes(b"current")
+            (copy_dir / "Ex 1_old.pdf").write_bytes(b"original")
+            (copy_dir / "Ex 2_new.pdf").write_bytes(b"generated")
+
+            self.assertEqual(module.main([str(evaluation), "--reset"]), 0)
+            self.assertFalse((evaluation / "correction.json").exists())
+            self.assertFalse((evaluation / "correction_progress.json").exists())
+            self.assertEqual((copy_dir / "Ex 1.pdf").read_bytes(), b"original")
+            self.assertFalse((copy_dir / "Ex 1_old.pdf").exists())
+            self.assertFalse((copy_dir / "Ex 2_new.pdf").exists())
+
+    def test_correction_batch_request_files_are_written_atomically(self) -> None:
+        module = self.modules["correction"]
+        with tempfile.TemporaryDirectory() as directory:
+            evaluation = Path(directory) / "Exam"
+            group_dir = evaluation / "Par label" / "Ex 1"
+            (evaluation / "Copies").mkdir(parents=True)
+            group_dir.mkdir(parents=True)
+            (evaluation / "labels").write_text("Ex 1\n", encoding="utf-8")
+            image = group_dir / "Group_1.jpg"
+            image.write_bytes(b"image")
+            atomic_write_json(
+                image.with_suffix(".json"),
+                [["01", 0, 400, 1.0, "Ex 1"]],
+            )
+            args = module.build_parser().parse_args([str(evaluation), "--batch"])
+            module.configure_runtime(
+                EvaluationWorkspace(evaluation),
+                [(str(image), "Ex 1")],
+                args,
+                api_client=Mock(),
+            )
+            with patch.object(module.prompting, "make_prompt", return_value="prompt"):
+                self.assertEqual(module.run_configured(args), 0)
+            lines = (evaluation / "batch_requests_flash.jsonl").read_text().splitlines()
+            self.assertEqual(len(lines), 1)
+            self.assertEqual(json.loads(lines[0])["key"], str(image))
+            self.assertEqual(
+                (evaluation / "batch_requests_pro.jsonl").read_text(), ""
+            )
+
+    def test_batch_helpers_use_mocked_api_and_atomic_combination(self) -> None:
+        submit = self.modules["submit_batches"]
+        fetch = self.modules["fetch_batched_results"]
+        status = self.modules["batch_status"]
+        with tempfile.TemporaryDirectory() as directory:
+            evaluation = Path(directory) / "Exam"
+            evaluation.mkdir()
+            (evaluation / "batch_requests_flash.jsonl").write_text(
+                "{}\n", encoding="utf-8"
+            )
+            client = Mock()
+            client.files.upload.return_value = SimpleNamespace(name="files/input")
+            client.batches.create.return_value = SimpleNamespace(name="batches/1")
+            self.assertEqual(
+                submit.run(EvaluationWorkspace(evaluation), client=client), 0
+            )
+            self.assertEqual(client.batches.create.call_count, 1)
+            manifest = json.loads((evaluation / "batch_jobs.json").read_text())
+            self.assertEqual(manifest["jobs"]["flash"]["name"], "batches/1")
+
+            jobs = [
+                SimpleNamespace(
+                    name="batches/1",
+                    display_name=f"flash-correction-{evaluation.name}",
+                    state=SimpleNamespace(name="JOB_STATE_SUCCEEDED"),
+                    dest=SimpleNamespace(file_name="files/flash-result"),
+                ),
+                SimpleNamespace(
+                    name="batches/2",
+                    display_name=f"pro-correction-{evaluation.name}",
+                    state=SimpleNamespace(name="JOB_STATE_SUCCEEDED"),
+                    dest=SimpleNamespace(file_name="files/pro-result"),
+                ),
+            ]
+            client.batches.get.return_value = jobs[0]
+            client.files.download.side_effect = [b'{"flash":1}\n', b'{"pro":1}']
+            self.assertEqual(
+                fetch.run(EvaluationWorkspace(evaluation), client=client), 0
+            )
+            self.assertEqual(
+                (evaluation / "batched_correction_result.jsonl").read_bytes(),
+                b'{"flash":1}\n',
+            )
+
+            job = SimpleNamespace(
+                state=SimpleNamespace(name="JOB_STATE_SUCCEEDED"),
+                dest=SimpleNamespace(file_name="files/result"),
+            )
+            client.batches.get.return_value = job
+            client.files.download.side_effect = None
+            client.files.download.return_value = b'{"downloaded":true}\n'
+            output = evaluation / "one-result.jsonl"
+            self.assertEqual(
+                status.download_job("batches/1", output=output, client=client), 0
+            )
+            self.assertEqual(output.read_bytes(), b'{"downloaded":true}\n')
+
     def test_post_correction_main_cleans_json_atomically(self) -> None:
         module = self.modules["post_correction"]
         with tempfile.TemporaryDirectory() as directory: