import io import tempfile import unittest from concurrent.futures import ThreadPoolExecutor, wait as real_wait from contextlib import redirect_stdout from pathlib import Path from threading import Event from unittest.mock import Mock, patch from copienator import CliError, EvaluationWorkspace, ExitCode from copienator.commands import gemini_for_labels as labels class LabelTimeoutTests(unittest.TestCase): def setUp(self): directory = tempfile.TemporaryDirectory() self.addCleanup(directory.cleanup) self.workspace = EvaluationWorkspace(Path(directory.name)) (self.workspace.root / "labels").write_text("Ex 1\n") (self.workspace.root / "names").write_text("Student\n") self.workspace.copies_dir.mkdir() (self.workspace.root / "Cutleft").mkdir() self.images = [] for name in ("Copie01_01", "Copie01_02", "Copie02_01"): image = self.workspace.root / "Cutleft" / f"{name}.jpg" image.write_bytes(b"image") self.images.append(image) self.client = Mock() self.response = Mock(text='{"name":"Student","list":[]}') def group(self, images): return labels.process_copy_group( self.workspace, "Copie01", images, client=self.client, labels_text="Ex 1\n", names_text="Student\n", valid_labels={"Ex 1"}, valid_names={"Student"}, overwrite=False, sleep=lambda _: None, target_interval=0, ) def test_cli_forwards_selected_concurrency(self): for arguments, expected in (([], 12), (["--workers", "2"], 2)): with self.subTest(arguments=arguments), \ patch.object(labels, "run", return_value=ExitCode.SUCCESS) as run: self.assertEqual(labels.main([str(self.workspace.root), *arguments]), 0) run.assert_called_once_with( unittest.mock.ANY, [self.workspace.root], overwrite=False, marked=False, max_workers=expected, ) def test_selected_concurrency_controls_executor(self): self.client.models.generate_content.return_value = self.response with patch.object(labels, "selected_images", return_value=(self.images, [])), \ patch.object(labels, "ThreadPoolExecutor", wraps=ThreadPoolExecutor) as executor: result = labels.run(self.workspace, [], client=self.client, sleep=lambda _: None, max_workers=1) self.assertEqual(result, ExitCode.SUCCESS) executor.assert_called_once_with(max_workers=1) self.assertEqual(self.client.models.generate_content.call_count, 3) def test_nonpositive_concurrency_rejected_before_requests(self): for workers in (0, -1): with self.subTest(workers=workers), self.assertRaises(CliError) as caught: labels.run(self.workspace, [], client=self.client, max_workers=workers) self.assertEqual(caught.exception.exit_code, ExitCode.INVALID_ARGUMENTS) self.client.models.generate_content.assert_not_called() def test_timeout_recovery_uses_bounded_requests_and_saves_result(self): self.client.models.generate_content.side_effect = [TimeoutError(), self.response] self.assertEqual(self.group(self.images[:1]), 1) self.assertEqual(self.client.models.generate_content.call_count, 2) for call in self.client.models.generate_content.call_args_list: options = call.kwargs["config"].http_options self.assertEqual(options.timeout, 80_000) self.assertEqual(options.retry_options.attempts, 1) self.assertTrue((self.workspace.copies_dir / "Copie01_01.json").exists()) def test_retries_use_low_thinking_and_next_image_restores_configured_level(self): self.client.models.generate_content.side_effect = [ TimeoutError(), TimeoutError(), self.response, self.response, ] with patch.object(labels.config, "THINKING_LEVEL_FOR_LABEL", "medium"): self.assertEqual(self.group(self.images[:2]), 2) self.assertEqual(labels.config.THINKING_LEVEL_FOR_LABEL, "medium") configs = [call.kwargs["config"] for call in self.client.models.generate_content.call_args_list] self.assertEqual( [cfg.thinking_config.thinking_level.value for cfg in configs], ["MEDIUM", "LOW", "LOW", "MEDIUM"], ) self.assertTrue(all(cfg.http_options.timeout == 80_000 for cfg in configs)) def test_exhaustion_defers_dependent_pages_but_other_copies_finish_and_resume(self): self.client.models.generate_content.side_effect = [ TimeoutError(), TimeoutError(), TimeoutError(), self.response, ] output = io.StringIO() with patch.object(labels, "selected_images", return_value=(self.images, [])), redirect_stdout(output): result = labels.run(self.workspace, [], client=self.client, sleep=lambda _: None, max_workers=1) self.assertEqual(result, ExitCode.PARTIAL) self.assertEqual(self.client.models.generate_content.call_count, 4) self.assertFalse((self.workspace.copies_dir / "Copie01_01.json").exists()) self.assertFalse((self.workspace.copies_dir / "Copie01_02.json").exists()) self.assertTrue((self.workspace.copies_dir / "Copie02_01.json").exists()) self.assertIn("Incomplete images", output.getvalue()) self.assertIn("Copie01_02.jpg", output.getvalue()) self.client.models.generate_content.reset_mock(side_effect=True) self.client.models.generate_content.return_value = self.response with patch.object(labels, "selected_images", return_value=(self.images, [])): result = labels.run(self.workspace, [], client=self.client, sleep=lambda _: None) self.assertEqual(result, ExitCode.SUCCESS) self.assertEqual(self.client.models.generate_content.call_count, 2) def test_changing_invalid_labels_cannot_retry_forever(self): self.client.models.generate_content.side_effect = [Mock(text=( '{"name":"Student","list":[{"box_2d":[1,2,3,4],"label":"bad%d"}]}' % index )) for index in range(3)] with patch.object(labels.config, "THINKING_LEVEL_FOR_LABEL", "medium"), \ self.assertRaises(labels.IncompleteCopyError): self.group(self.images[:1]) self.assertEqual(self.client.models.generate_content.call_count, 3) self.assertEqual( [call.kwargs["config"].thinking_config.thinking_level.value for call in self.client.models.generate_content.call_args_list], ["MEDIUM", "LOW", "LOW"], ) def test_progress_reports_waiting_image_and_elapsed_time(self): entered, release = Event(), Event() self.addCleanup(release.set) def request(**kwargs): entered.set() if not release.wait(2): raise AssertionError("Test did not release request") return self.response first_wait = True def wait(futures, timeout): nonlocal first_wait self.assertEqual(timeout, 30) if first_wait: first_wait = False self.assertTrue(entered.wait(2)) return set(), set(futures) release.set() return real_wait(futures, timeout=2) self.client.models.generate_content.side_effect = request output = io.StringIO() with patch.object(labels, "selected_images", return_value=(self.images[:1], [])), \ patch.object(labels, "wait", side_effect=wait), redirect_stdout(output): result = labels.run(self.workspace, [], client=self.client, sleep=lambda _: None) self.assertEqual(result, ExitCode.SUCCESS) self.assertIn("Label detection progress: 0/1 copy groups completed.", output.getvalue()) self.assertIn("Label detection progress: 1/1 copy groups completed.", output.getvalue()) self.assertIn("1 copies outstanding", output.getvalue()) self.assertIn("Copie01_01.jpg, attempt 1/3", output.getvalue()) self.assertIn("elapsed for image", output.getvalue())