182 lines
9.4 KiB
Python
182 lines
9.4 KiB
Python
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
|
|
from copienator.copy_errors import mark_copy_error
|
|
|
|
|
|
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_flagged_reruns_get_fresh_seeds_and_replace_existing_results(self):
|
|
pdf = self.workspace.copies_dir / "Copie01.pdf"
|
|
pdf.touch()
|
|
mark_copy_error(self.workspace, pdf, "Review", source="plotting")
|
|
self.client.models.generate_content.return_value = self.response
|
|
with patch.object(labels, "selected_images", return_value=(self.images[:1], [])), \
|
|
patch.object(labels, "fresh_seed", side_effect=[101, 202]), \
|
|
patch.object(labels.config, "THINKING_LEVEL_FOR_LABEL", "medium"):
|
|
for _ in range(2):
|
|
self.assertEqual(labels.run(
|
|
self.workspace, [], marked=True, client=self.client,
|
|
sleep=lambda _: None, max_workers=1,
|
|
), ExitCode.SUCCESS)
|
|
configs = [call.kwargs["config"]
|
|
for call in self.client.models.generate_content.call_args_list]
|
|
self.assertEqual([cfg.seed for cfg in configs], [101, 202])
|
|
self.assertEqual([cfg.thinking_config.thinking_level.value for cfg in configs],
|
|
["MEDIUM", "MEDIUM"])
|
|
|
|
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))
|
|
self.assertNotEqual(configs[0].seed, configs[1].seed)
|
|
self.assertNotEqual(configs[1].seed, configs[2].seed)
|
|
|
|
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())
|