"""Regression tests for the RunPod helper scripts.""" from __future__ import annotations import importlib.util import json import os import subprocess import sys import time from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[2] def _load_module(module_name: str, relative_path: str): module_path = REPO_ROOT / relative_path script_dir = str(module_path.parent) added_path = False if script_dir not in sys.path: sys.path.insert(0, script_dir) added_path = True spec = importlib.util.spec_from_file_location(module_name, module_path) module = importlib.util.module_from_spec(spec) assert spec.loader is not None try: spec.loader.exec_module(module) finally: if added_path: sys.path.remove(script_dir) return module def test_process_local_videos_times_out_and_continues(tmp_path, monkeypatch): module = _load_module("test_process_local_videos_script", "scripts/process-local-videos.py") class _FakeImageSearch: def get_processed_video_keys(self): return set() def is_video_processed(self, natural_key): return False def _fake_process_video(natural_key, label, delete_video_after): if natural_key == "video-1": time.sleep(2) return {"status": "success", "message": f"processed {natural_key}"} monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr(module, "get_all_local_video_keys", lambda: ["video-1", "video-2"]) monkeypatch.setattr(module, "process_video", _fake_process_video) report_path = tmp_path / "timeout-report.json" monkeypatch.setattr( sys, "argv", [ "process-local-videos.py", "--timeout-seconds", "1", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 1 report = json.loads(report_path.read_text()) assert report["timeout_seconds"] == 1 assert report["processed"] == 1 assert report["failed"] == 1 assert report["results"][0]["natural_key"] == "video-1" assert report["results"][0]["status"] == "error" assert "Timed out after 1s" in report["results"][0]["message"] assert report["results"][1]["natural_key"] == "video-2" assert report["results"][1]["status"] == "success" def test_process_local_videos_can_pin_gpu(monkeypatch): module = _load_module("test_process_local_videos_gpu_pin", "scripts/process-local-videos.py") monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) monkeypatch.setattr( sys, "argv", [ "process-local-videos.py", "--gpu-id", "1", "--limit", "0", ], ) captured = {} def _fake_load_backend_dependencies(): captured["gpu_id"] = os.environ.get("CUDA_VISIBLE_DEVICES") module.ImageSearch = lambda: type("_Fake", (), {"get_processed_video_keys": lambda self: set()})() module.get_all_local_video_keys = lambda: [] module.process_video = lambda *args, **kwargs: {"status": "success", "message": "ok"} module.atomic_write_json = lambda *args, **kwargs: None monkeypatch.setattr(module, "_load_backend_dependencies", _fake_load_backend_dependencies) exit_code = module.main() assert exit_code == 0 assert captured["gpu_id"] == "1" assert os.environ["CUDA_VISIBLE_DEVICES"] == "1" def test_process_local_videos_report_write_is_best_effort(tmp_path): module = _load_module("test_process_local_videos_report_best_effort", "scripts/process-local-videos.py") calls = [] def _fake_atomic_write(path, payload, indent=None): calls.append(path) if "search-ui-report-shadow" in path: Path(path).parent.mkdir(parents=True, exist_ok=True) Path(path).write_text(json.dumps(payload)) return raise OSError(5, "Input/output error") module.atomic_write_json = _fake_atomic_write module._save_report(str(tmp_path / "report.json"), {"status": "ok"}) assert any("search-ui-report-shadow" in path for path in calls) def test_process_local_videos_filters_partition_before_sharding(monkeypatch): module = _load_module("test_process_local_videos_partition", "scripts/process-local-videos.py") partition_module = _load_module("test_batch_partition_helper", "backend/batch_partition.py") class _FakeImageSearch: def get_processed_video_keys(self): return set() def is_video_processed(self, natural_key): return False processed = [] monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr( module, "get_all_local_video_keys", lambda: [f"video-{idx}" for idx in range(12)], ) monkeypatch.setattr( module, "process_video", lambda natural_key, label, delete_video_after: processed.append(natural_key) or {"status": "success", "message": natural_key}, ) module.atomic_write_json = lambda *args, **kwargs: None monkeypatch.setattr( sys, "argv", [ "process-local-videos.py", "--partition-total", "2", "--partition-index", "1", "--shard-count", "2", "--shard-index", "0", ], ) exit_code = module.main() assert exit_code == 0 expected_partition = partition_module.filter_natural_keys_for_partition( [f"video-{idx}" for idx in range(12)], partition_total=2, partition_index=1, ) expected_shard = [ natural_key for offset, natural_key in enumerate(expected_partition) if offset % 2 == 0 ] assert processed == expected_shard def test_process_local_videos_global_limit_is_exact_across_shards(monkeypatch): module = _load_module("test_process_local_videos_exact_limit", "scripts/process-local-videos.py") class _FakeImageSearch: def get_processed_video_keys(self): return set() def is_video_processed(self, natural_key): return False processed = [] monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr(module, "get_all_local_video_keys", lambda: [f"video-{idx}" for idx in range(10)]) monkeypatch.setattr( module, "process_video", lambda natural_key, label, delete_video_after: processed.append(natural_key) or {"status": "success", "message": natural_key}, ) module.atomic_write_json = lambda *args, **kwargs: None for shard_index in (0, 1): monkeypatch.setattr( sys, "argv", [ "process-local-videos.py", "--limit", "5", "--shard-count", "2", "--shard-index", str(shard_index), ], ) exit_code = module.main() assert exit_code == 0 assert sorted(processed) == [f"video-{idx}" for idx in range(5)] def test_process_local_videos_target_file_filters_and_preserves_order(tmp_path, monkeypatch): module = _load_module("test_process_local_videos_target_file", "scripts/process-local-videos.py") class _FakeImageSearch: def get_processed_video_keys(self): return {"video-2"} def is_video_processed(self, natural_key): return natural_key == "video-2" target_file = tmp_path / "targets.txt" target_file.write_text("video-3\nvideo-1\nvideo-2\nvideo-99\n", encoding="utf-8") processed = [] monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr(module, "get_all_local_video_keys", lambda: ["video-1", "video-2", "video-3", "video-4"]) monkeypatch.setattr( module, "process_video", lambda natural_key, label, delete_video_after: processed.append(natural_key) or {"status": "success", "message": natural_key}, ) module.atomic_write_json = lambda *args, **kwargs: None monkeypatch.setattr( sys, "argv", [ "process-local-videos.py", "--natural-key-file", str(target_file), ], ) exit_code = module.main() assert exit_code == 0 assert processed == ["video-3", "video-1"] def test_download_all_videos_removes_stale_part_files(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_script", "scripts/download-all-videos.py") stale_dir = tmp_path / "video-1" / "720p" stale_dir.mkdir(parents=True) stale_part = stale_dir / "video.mp4.123.part" stale_part.write_bytes(b"partial") monkeypatch.setattr(module, "get_videos_dir", lambda: str(tmp_path)) monkeypatch.setattr(module, "refresh_catalog", lambda language, redownload=False: {}) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) report_path = tmp_path / "download-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 0 assert not stale_part.exists() report = json.loads(report_path.read_text()) assert report["partial_files_removed"] == 1 assert report["total"] == 0 def test_download_all_videos_report_write_is_best_effort(tmp_path): module = _load_module("test_download_all_videos_report_best_effort", "scripts/download-all-videos.py") calls = [] def _fake_atomic_write(path, payload, indent=None): calls.append(path) if "search-ui-report-shadow" in path: Path(path).parent.mkdir(parents=True, exist_ok=True) Path(path).write_text(json.dumps(payload)) return raise OSError(5, "Input/output error") module.atomic_write_json = _fake_atomic_write module._save_report(str(tmp_path / "download-report.json"), {"status": "ok"}) assert any("search-ui-report-shadow" in path for path in calls) def test_download_all_videos_skips_processed_keys_by_default(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_skip_processed", "scripts/download-all-videos.py") monkeypatch.setattr( module, "refresh_catalog", lambda language, redownload=False: { "video-1": {"languageAgnosticNaturalKey": "video-1", "title": "Processed"}, "video-2": {"languageAgnosticNaturalKey": "video-2", "title": "Pending"}, }, ) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) monkeypatch.setattr(module, "find_video_file", lambda media_item, label, subtitled=False: {"label": label}) monkeypatch.setattr(module, "download_video", lambda media_item, label: f"/tmp/{media_item['languageAgnosticNaturalKey']}.mp4") class _FakeImageSearch: def get_processed_video_keys(self): return {"video-1"} monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) report_path = tmp_path / "download-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 0 report = json.loads(report_path.read_text()) assert report["total"] == 1 assert report["processed"] == 1 assert report["downloaded"] == 1 assert report["skipped_already_processed"] == 1 assert report["results"][0]["natural_key"] == "video-2" def test_download_all_videos_applies_partition_before_limit(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_partition", "scripts/download-all-videos.py") catalog = { f"video-{idx}": {"languageAgnosticNaturalKey": f"video-{idx}", "title": f"Video {idx}"} for idx in range(8) } monkeypatch.setattr(module, "refresh_catalog", lambda language, redownload=False: catalog) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) monkeypatch.setattr(module, "find_video_file", lambda media_item, label, subtitled=False: {"label": label}) monkeypatch.setattr(module, "download_video", lambda media_item, label: f"/tmp/{media_item['languageAgnosticNaturalKey']}.mp4") class _FakeImageSearch: def get_processed_video_keys(self): return set() monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) report_path = tmp_path / "download-partition-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--partition-total", "2", "--partition-index", "0", "--limit", "2", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 0 report = json.loads(report_path.read_text()) assert report["partition_total"] == 2 assert report["partition_index"] == 0 assert report["processed"] == 2 def test_download_all_videos_pushes_monthly_programs_to_back(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_priority", "scripts/download-all-videos.py") monkeypatch.setattr( module, "refresh_catalog", lambda language, redownload=False: { "pub-jwb-135_1_VIDEO": { "languageAgnosticNaturalKey": "pub-jwb-135_1_VIDEO", "title": "JW Broadcasting—April 2026", "_category_key": "VODStudio", "_subcategory": "Monthly Programs", }, "talk-1": { "languageAgnosticNaturalKey": "talk-1", "title": "Talk 1", "_category_key": "VODStudio", "_subcategory": "Apply Bible Principles", }, "insert-1": { "languageAgnosticNaturalKey": "insert-1", "title": "Insert 1", "_category_key": "VODStudio", "_subcategory": "Become Jehovah's Friend", }, }, ) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) monkeypatch.setattr(module, "find_video_file", lambda media_item, label, subtitled=False: {"label": label}) monkeypatch.setattr(module, "download_video", lambda media_item, label: f"/tmp/{media_item['languageAgnosticNaturalKey']}.mp4") class _FakeImageSearch: def get_processed_video_keys(self): return set() monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) report_path = tmp_path / "download-priority-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 0 report = json.loads(report_path.read_text()) assert [entry["natural_key"] for entry in report["results"]] == [ "insert-1", "talk-1", "pub-jwb-135_1_VIDEO", ] def test_download_all_videos_plan_only_writes_fixed_target_file(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_plan_only", "scripts/download-all-videos.py") catalog = { f"video-{idx}": {"languageAgnosticNaturalKey": f"video-{idx}", "title": f"Video {idx}"} for idx in range(6) } monkeypatch.setattr(module, "refresh_catalog", lambda language, redownload=False: catalog) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) class _FakeImageSearch: def get_processed_video_keys(self): return {"video-0"} monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr(module, "_cleanup_partial_downloads", lambda: 99) target_path = tmp_path / "targets.txt" report_path = tmp_path / "download-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--partition-total", "2", "--partition-index", "1", "--limit", "2", "--write-target-file", str(target_path), "--plan-only", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 0 planned = [line.strip() for line in target_path.read_text().splitlines() if line.strip()] assert len(planned) == 2 report = json.loads(report_path.read_text()) assert report["total"] == 2 assert report["processed"] == 0 assert report["partial_files_removed"] == 0 def test_download_all_videos_natural_key_file_preserves_fixed_chunk(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_target_file", "scripts/download-all-videos.py") catalog = { f"video-{idx}": {"languageAgnosticNaturalKey": f"video-{idx}", "title": f"Video {idx}"} for idx in range(4) } monkeypatch.setattr(module, "refresh_catalog", lambda language, redownload=False: catalog) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) monkeypatch.setattr(module, "find_video_file", lambda media_item, label, subtitled=False: {"label": label}) monkeypatch.setattr(module, "download_video", lambda media_item, label: f"/tmp/{media_item['languageAgnosticNaturalKey']}.mp4") class _FakeImageSearch: def get_processed_video_keys(self): return {"video-2"} monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) target_path = tmp_path / "targets.txt" target_path.write_text("video-3\nvideo-2\nvideo-1\n", encoding="utf-8") report_path = tmp_path / "download-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--natural-key-file", str(target_path), "--limit", "1", "--partition-total", "2", "--partition-index", "0", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 0 report = json.loads(report_path.read_text()) assert report["total"] == 2 assert {result["natural_key"] for result in report["results"]} == {"video-3", "video-1"} def test_download_all_videos_plan_only_refilters_target_file_against_processed_db(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_target_refresh", "scripts/download-all-videos.py") catalog = { f"video-{idx}": {"languageAgnosticNaturalKey": f"video-{idx}", "title": f"Video {idx}"} for idx in range(4) } monkeypatch.setattr(module, "refresh_catalog", lambda language, redownload=False: catalog) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) class _FakeImageSearch: def get_processed_video_keys(self): return {"video-2"} monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr(module, "_cleanup_partial_downloads", lambda: 0) source_path = tmp_path / "targets.txt" source_path.write_text("video-3\nvideo-2\nvideo-1\n", encoding="utf-8") refreshed_path = tmp_path / "targets-refreshed.txt" report_path = tmp_path / "download-plan-refresh-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--natural-key-file", str(source_path), "--write-target-file", str(refreshed_path), "--plan-only", "--report", str(report_path), ], ) exit_code = module.main() assert exit_code == 0 assert refreshed_path.read_text(encoding="utf-8").splitlines() == ["video-3", "video-1"] report = json.loads(report_path.read_text()) assert report["total"] == 2 assert report["processed"] == 0 def test_download_all_videos_retries_transient_failures(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_retry", "scripts/download-all-videos.py") monkeypatch.setattr( module, "refresh_catalog", lambda language, redownload=False: { "video-1": {"languageAgnosticNaturalKey": "video-1", "title": "Retry me"}, }, ) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) monkeypatch.setattr(module, "find_video_file", lambda media_item, label, subtitled=False: {"label": label}) attempts = {"count": 0} def _fake_download_video(media_item, label): attempts["count"] += 1 if attempts["count"] == 1: return None return f"/tmp/{media_item['languageAgnosticNaturalKey']}.mp4" class _FakeImageSearch: def get_processed_video_keys(self): return set() monkeypatch.setattr(module, "download_video", _fake_download_video) monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr(module.time, "sleep", lambda *_args, **_kwargs: None) report_path = tmp_path / "download-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--report", str(report_path), "--download-retries", "3", "--retry-backoff-seconds", "0", ], ) exit_code = module.main() assert exit_code == 0 assert attempts["count"] == 2 report = json.loads(report_path.read_text()) assert report["downloaded"] == 1 assert report["errors"] == 0 assert report["results"][0]["status"] == "downloaded" def test_download_all_videos_records_final_item_error_without_crashing(tmp_path, monkeypatch): module = _load_module("test_download_all_videos_item_error", "scripts/download-all-videos.py") monkeypatch.setattr( module, "refresh_catalog", lambda language, redownload=False: { "video-1": {"languageAgnosticNaturalKey": "video-1", "title": "Broken"}, }, ) monkeypatch.setattr(module, "get_existing_video_keys", lambda label: set()) monkeypatch.setattr(module, "find_video_file", lambda media_item, label, subtitled=False: {"label": label}) monkeypatch.setattr(module, "download_video", lambda media_item, label: None) class _FakeImageSearch: def get_processed_video_keys(self): return set() monkeypatch.setattr(module, "ImageSearch", lambda: _FakeImageSearch()) monkeypatch.setattr(module.time, "sleep", lambda *_args, **_kwargs: None) report_path = tmp_path / "download-report.json" monkeypatch.setattr( sys, "argv", [ "download-all-videos.py", "--report", str(report_path), "--download-retries", "2", "--retry-backoff-seconds", "0", ], ) exit_code = module.main() assert exit_code == 1 report = json.loads(report_path.read_text()) assert report["processed"] == 1 assert report["errors"] == 1 assert report["results"][0]["status"] == "error" assert "after 2 attempts" in report["results"][0]["message"] def test_evaluate_runpod_output_returns_nonzero_when_checks_fail(monkeypatch): module = _load_module("test_evaluate_runpod_output_script", "scripts/evaluate-runpod-output.py") monkeypatch.setattr(module, "_collect_phase_reports", lambda report_dir: {}) monkeypatch.setattr( module, "_collect_filesystem", lambda language: { "subtitles_root": "/tmp/subtitles", "videos_root": "/tmp/videos", "vtt_files": 0, "txt_files": 0, "mp4_files": 0, "thumbnails": 0, "small_thumbnails": 0, "face_crops": 0, }, ) monkeypatch.setattr(module, "_gather_search_samples", lambda language: {"stats": {"total_subtitles": 0}}) monkeypatch.setattr(module, "_gather_image_samples", lambda: {"stats": {"unique_videos": 0}}) monkeypatch.setattr(module, "_gather_face_samples", lambda: {"stats": {"total_faces": 0}}) monkeypatch.setattr(sys, "argv", ["evaluate-runpod-output.py"]) exit_code = module.main() assert exit_code == 1 def test_evaluate_runpod_output_counts_face_crops_directory(tmp_path, monkeypatch): module = _load_module("test_evaluate_runpod_output_face_crops", "scripts/evaluate-runpod-output.py") subtitles_root = tmp_path / "subtitles" videos_root = tmp_path / "videos" (subtitles_root / "E").mkdir(parents=True) face_crop_dir = videos_root / "video-1" / "720p" / "face-crops" face_crop_dir.mkdir(parents=True) (face_crop_dir / "frame1_face0.jpg").write_bytes(b"jpeg") monkeypatch.setattr(module, "get_subtitles_dir", lambda: str(subtitles_root)) monkeypatch.setattr(module, "get_videos_dir", lambda: str(videos_root)) filesystem = module._collect_filesystem("E") assert filesystem["face_crops"] == 1 def test_runpod_batch_worker_splits_process_limit_across_shards(tmp_path, monkeypatch): module = _load_module("test_runpod_batch_worker_script", "scripts/runpod-batch-worker.py") class _FakeProc: def __init__(self, cmd, **kwargs): self.cmd = cmd def poll(self): return 0 def wait(self): return 0 def terminate(self): return None def kill(self): return None monkeypatch.setattr(module.subprocess, "Popen", lambda cmd, **kwargs: _FakeProc(cmd, **kwargs)) monkeypatch.setattr(module.time, "sleep", lambda *_args, **_kwargs: None) manifest_path = tmp_path / "manifest.json" manifest = {"manifest_path": str(manifest_path), "phases": []} report_dir = tmp_path / "reports" report_dir.mkdir() module._run_processing_shards( label="720p", process_shards=2, process_limit=5, partition_total=1, partition_index=0, delete_after=False, report_dir=str(report_dir), manifest=manifest, shard_timeout_seconds=0, ) commands = [shard["command"] for shard in manifest["phases"][0]["shards"]] limits = [cmd[cmd.index("--limit") + 1] for cmd in commands] assert limits == ["5", "5"] def test_runpod_batch_worker_passes_gpu_ids_to_shards(tmp_path, monkeypatch): module = _load_module("test_runpod_batch_worker_gpu_ids", "scripts/runpod-batch-worker.py") class _FakeProc: def __init__(self, cmd, **kwargs): self.cmd = cmd def poll(self): return 0 def wait(self): return 0 def terminate(self): return None def kill(self): return None monkeypatch.setattr(module.subprocess, "Popen", lambda cmd, **kwargs: _FakeProc(cmd, **kwargs)) monkeypatch.setattr(module.time, "sleep", lambda *_args, **_kwargs: None) monkeypatch.setenv("RUNPOD_WORKER_GPU_IDS", "3,5") manifest_path = tmp_path / "manifest.json" manifest = {"manifest_path": str(manifest_path), "phases": []} report_dir = tmp_path / "reports" report_dir.mkdir() module._run_processing_shards( label="720p", process_shards=2, process_limit=0, partition_total=1, partition_index=0, delete_after=False, report_dir=str(report_dir), manifest=manifest, shard_timeout_seconds=0, ) commands = [shard["command"] for shard in manifest["phases"][0]["shards"]] gpu_ids = [cmd[cmd.index("--gpu-id") + 1] for cmd in commands] assert gpu_ids == ["3", "5"] def test_runpod_batch_worker_respects_single_visible_cuda_device(tmp_path, monkeypatch): module = _load_module("test_runpod_batch_worker_single_visible", "scripts/runpod-batch-worker.py") class _FakeProc: def __init__(self, cmd, **kwargs): self.cmd = cmd def poll(self): return 0 def wait(self): return 0 def terminate(self): return None def kill(self): return None monkeypatch.setattr(module.subprocess, "Popen", lambda cmd, **kwargs: _FakeProc(cmd, **kwargs)) monkeypatch.setattr(module.time, "sleep", lambda *_args, **_kwargs: None) monkeypatch.delenv("RUNPOD_WORKER_GPU_IDS", raising=False) monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7") manifest_path = tmp_path / "manifest.json" manifest = {"manifest_path": str(manifest_path), "phases": []} report_dir = tmp_path / "reports" report_dir.mkdir() module._run_processing_shards( label="720p", process_shards=1, process_limit=0, partition_total=1, partition_index=0, delete_after=False, report_dir=str(report_dir), manifest=manifest, shard_timeout_seconds=0, ) command = manifest["phases"][0]["shards"][0]["command"] assert command[command.index("--gpu-id") + 1] == "7" def test_runpod_batch_worker_passes_partition_args_to_shards(tmp_path, monkeypatch): module = _load_module("test_runpod_batch_worker_partition", "scripts/runpod-batch-worker.py") class _FakeProc: def __init__(self, cmd, **kwargs): self.cmd = cmd def poll(self): return 0 def wait(self): return 0 def terminate(self): return None def kill(self): return None monkeypatch.setattr(module.subprocess, "Popen", lambda cmd, **kwargs: _FakeProc(cmd, **kwargs)) monkeypatch.setattr(module.time, "sleep", lambda *_args, **_kwargs: None) manifest_path = tmp_path / "manifest.json" manifest = {"manifest_path": str(manifest_path), "phases": []} report_dir = tmp_path / "reports" report_dir.mkdir() module._run_processing_shards( label="720p", process_shards=2, process_limit=0, partition_total=3, partition_index=1, delete_after=False, report_dir=str(report_dir), manifest=manifest, shard_timeout_seconds=0, ) commands = [shard["command"] for shard in manifest["phases"][0]["shards"]] assert all("--partition-total" in cmd for cmd in commands) assert all("--partition-index" in cmd for cmd in commands) def test_run_batch_script_wires_gpu_ids_into_process_workers(): script = (REPO_ROOT / "deploy/runpod/run-batch.sh").read_text() multi_gpu_proof = (REPO_ROOT / "deploy/runpod/run-multi-gpu-proof.sh").read_text() assert "RUNPOD_GPU_IDS" in script assert "get_shard_gpu_id" in script assert 'shard_cmd+=(--gpu-id "$gpu_id")' in script assert "RUNPOD_PARTITION_TOTAL" in script assert "--partition-total \"$RUNPOD_PARTITION_TOTAL\"" in script assert "count_natural_key_file_unprocessed" in script assert "DOWNLOAD_TARGET_FILE" in script assert "DOWNLOAD_PLAN_FILE" in script assert "DOWNLOAD_PLANNED_TARGET_COUNT" in script assert "restart_download_if_needed" in script assert '--natural-key-file "$DOWNLOAD_TARGET_FILE"' in script assert "restart_dead_shards_if_needed" in script assert "WATCHDOG_STATE_DIR" in script assert "/tmp/search-ui-watchdog" in script assert "SHARD_DISABLED" in script assert "backup_runtime_databases_best_effort" in script assert "RunPod Multi-GPU Proof" in multi_gpu_proof assert "export RUNPOD_GPU_IDS" in multi_gpu_proof assert "duplicate GPU ID" in multi_gpu_proof assert "RUNPOD_PARTITION_TOTAL and RUNPOD_PARTITION_INDEX must be set explicitly" in multi_gpu_proof assert 'bash "$SCRIPT_DIR/run-batch.sh"' in multi_gpu_proof def test_run_batch_archives_shard_progress_before_restart_cap_failure(): script = (REPO_ROOT / "deploy/runpod/run-batch.sh").read_text() archive_idx = script.index(' archive_shard_progress "$shard_index"') cap_idx = script.index(' if [ "${SHARD_RESTART_COUNTS[$shard_index]:-0}" -ge "$WATCHDOG_MAX_SHARD_RESTARTS" ]; then') assert archive_idx < cap_idx def test_run_proof_validation_uses_vec_connections(): script = (REPO_ROOT / "deploy/runpod/run-proof.sh").read_text() assert "from db_utils import IMAGE_DB_PRAGMAS, connect_vec_db" in script assert "image_conn = connect_vec_db(get_image_db_path(), pragmas=IMAGE_DB_PRAGMAS)" in script assert "connect_vec_db(face_db_path, pragmas=IMAGE_DB_PRAGMAS)" in script def test_runpod_scripts_fail_fast_when_system_tools_are_missing(): run_proof = (REPO_ROOT / "deploy/runpod/run-proof.sh").read_text() run_batch = (REPO_ROOT / "deploy/runpod/run-batch.sh").read_text() export_results = (REPO_ROOT / "deploy/runpod/export-results.sh").read_text() setup_pod = (REPO_ROOT / "deploy/runpod/setup-pod.sh").read_text() watchdog = (REPO_ROOT / "deploy/runpod/watchdog.py").read_text() check_gpu_runtime = (REPO_ROOT / "deploy/runpod/check_gpu_runtime.py").read_text() backup_runtime_dbs = (REPO_ROOT / "deploy/runpod/backup_runtime_dbs.py").read_text() restore_runtime_dbs = (REPO_ROOT / "deploy/runpod/restore_runtime_dbs.py").read_text() assert "require_system_commands ffmpeg ffprobe sqlite3" in run_proof assert "Recovery: bash deploy/runpod/setup-pod.sh" in run_proof assert "require_system_commands ffmpeg ffprobe sqlite3 setsid" in run_batch assert "check_gpu_runtime" in run_batch assert 'check_gpu_runtime.py" --context "run-proof.sh"' in run_proof assert 'check_gpu_runtime.py" --context "run-batch.sh"' in run_batch assert 'check_gpu_runtime.py" --context "setup-pod.sh"' in setup_pod assert 'backup_runtime_dbs.py' in run_proof assert 'backup_runtime_dbs.py' in run_batch assert "backup_runtime_databases_or_die" in run_batch assert "count_natural_key_file_processed_in_db" in run_batch assert "All downloadable targets are already processed in the DB." in run_batch assert "Recovery: bash deploy/runpod/setup-pod.sh" in run_batch assert "sqlite3 unavailable" in export_results assert "Recovery: bash deploy/runpod/setup-pod.sh" in export_results assert "db-backup (runtime backup failed)" in export_results assert "--include-large-thumbnails" in export_results assert "excluding large thumbnails (default)" in export_results assert 'restore_runtime_dbs.py' in setup_pod assert 'VENV_DIR=/tmp/search-ui-venv' in setup_pod assert 'run_pip_with_retries' in setup_pod assert "--auto-export" in run_batch assert "Allowing active shards to finish the downloaded backlog" in run_batch assert "downloader exited too many times" in run_batch assert "CHUNK COMPLETE WITH DOWNLOAD WARNINGS" in run_batch assert "all shards are disabled while downloader is still running" in run_batch assert "watchdog.py" in run_batch assert "write_watchdog_state" in run_batch assert "cleanup_on_exit" in run_batch assert "trap 'cleanup_on_exit $?'" in run_batch assert 'setsid "${download_cmd[@]}"' in run_batch assert "os.killpg(pid, signal.SIGTERM)" in watchdog assert "--dead-shard-checks" in watchdog assert "GPU capability" in check_gpu_runtime assert "CUDA smoke test failed" in check_gpu_runtime assert "source_conn.backup(target_conn)" in backup_runtime_dbs assert "backup newer than runtime" in restore_runtime_dbs assert "backup-manifest.json" in backup_runtime_dbs assert "manifest_sha_mismatch" in restore_runtime_dbs assert "export-status.json" in export_results assert "_collect_export_status" in watchdog assert "RUNPOD_ENABLE_DEEPFACE_GPU" in setup_pod assert "tensorflow[and-cuda]" in setup_pod assert "Linked ptxas" in setup_pod assert "SEARCH_UI_TF_LIBRARY_DIRS" in setup_pod assert "LD_LIBRARY_PATH" in setup_pod assert "RUNPOD_PARTITION_TOTAL" in run_batch def test_check_gpu_runtime_rejects_unsupported_arch(): module = _load_module("test_runpod_check_gpu_runtime", "deploy/runpod/check_gpu_runtime.py") probe = { "torch_version": "2.4.1+cu124", "cuda_version": "12.4", "cuda_available": True, "device_count": 1, "device_name": "NVIDIA RTX PRO 4500 Blackwell", "total_memory_gb": 32.0, "capability_tag": "sm_120", "supported_arches": ["sm_50", "sm_60", "sm_70", "sm_75", "sm_80", "sm_86", "sm_90"], "smoke_error": None, } ok, message = module.evaluate_runtime(gpu_visible=True, probe=probe) assert ok is False assert "sm_120" in message assert "sm_90" in message def test_watchdog_terminates_orphaned_children(tmp_path): module = _load_module("test_runpod_watchdog", "deploy/runpod/watchdog.py") download_proc = subprocess.Popen(["sleep", "300"]) shard_proc = subprocess.Popen(["sleep", "300"]) try: state_path = tmp_path / "watchdog-state.json" status_path = tmp_path / "watchdog-status.json" state_path.write_text( json.dumps( { "phase": "processing", "heartbeat_epoch": int(time.time()), "supervisor_pid": 999999, "download_pid": download_proc.pid, "shard_pids": [shard_proc.pid], "report_dir": str(tmp_path), "workspace_dir": str(tmp_path), "download_finished": False, "download_exit_code": 0, } ) ) exit_code = module.main( [ "--state-file", str(state_path), "--status-file", str(status_path), "--once", ] ) assert exit_code == 2 download_proc.wait(timeout=10) shard_proc.wait(timeout=10) status = json.loads(status_path.read_text()) assert status["action"] == "terminated" assert status["reason"] == "supervisor_dead" finally: for proc in (download_proc, shard_proc): if proc.poll() is None: proc.terminate() try: proc.wait(timeout=5) except subprocess.TimeoutExpired: proc.kill() def test_watchdog_requests_restart_on_partial_shard_death(tmp_path): module = _load_module("test_runpod_watchdog_partial", "deploy/runpod/watchdog.py") alive_shard = subprocess.Popen(["sleep", "300"]) dead_shard = subprocess.Popen(["sleep", "300"]) dead_shard.terminate() dead_shard.wait(timeout=5) terminated = {} try: state_path = tmp_path / "watchdog-state.json" status_path = tmp_path / "watchdog-status.json" state_path.write_text( json.dumps( { "phase": "processing", "heartbeat_epoch": int(time.time()), "supervisor_pid": os.getpid(), "download_pid": None, "shard_pids": [alive_shard.pid, dead_shard.pid], "report_dir": str(tmp_path), "workspace_dir": str(tmp_path), "download_finished": True, "download_exit_code": 0, } ) ) module._terminate_pids = lambda pids, grace_seconds=5.0: terminated.setdefault("pids", list(pids)) exit_code = module.main( [ "--state-file", str(state_path), "--status-file", str(status_path), "--dead-shard-checks", "1", "--once", ] ) assert exit_code == 0 status = json.loads(status_path.read_text()) assert status["action"] == "restart_requested" assert status["reason"] == "partial_shard_death" assert dead_shard.pid in status["dead_shard_pids"] assert "pids" not in terminated finally: if alive_shard.poll() is None: alive_shard.terminate() try: alive_shard.wait(timeout=5) except subprocess.TimeoutExpired: alive_shard.kill() def test_watchdog_uses_in_memory_state_when_status_writes_fail(tmp_path, monkeypatch): module = _load_module("test_runpod_watchdog_memory_state", "deploy/runpod/watchdog.py") alive_shard = subprocess.Popen(["sleep", "300"]) dead_shard = subprocess.Popen(["sleep", "300"]) dead_shard.terminate() dead_shard.wait(timeout=5) try: state_path = tmp_path / "watchdog-state.json" status_path = tmp_path / "watchdog-status.json" state_path.write_text( json.dumps( { "phase": "processing", "heartbeat_epoch": int(time.time()), "supervisor_pid": os.getpid(), "download_pid": None, "shard_pids": [alive_shard.pid, dead_shard.pid], "report_dir": str(tmp_path), "workspace_dir": str(tmp_path), "download_finished": True, "download_exit_code": 0, } ) ) monkeypatch.setattr(module, "_save_status", lambda path, payload: None) exit_code, previous = module._evaluate( module.parse_args( [ "--state-file", str(state_path), "--status-file", str(status_path), "--dead-shard-checks", "2", "--once", ] ), str(status_path), {}, ) assert exit_code == 0 assert previous["dead_shard_streak"] == 1 exit_code, previous = module._evaluate( module.parse_args( [ "--state-file", str(state_path), "--status-file", str(status_path), "--dead-shard-checks", "2", "--once", ] ), str(status_path), previous, ) assert exit_code == 0 assert previous["action"] == "restart_requested" assert previous["dead_shard_streak"] == 2 finally: if alive_shard.poll() is None: alive_shard.terminate() try: alive_shard.wait(timeout=5) except subprocess.TimeoutExpired: alive_shard.kill() def test_watchdog_terminates_stale_pipeline_when_state_file_disappears(tmp_path): module = _load_module("test_runpod_watchdog_missing_state", "deploy/runpod/watchdog.py") status_path = tmp_path / "watchdog-status.json" previous = { "action": "healthy", "phase": "processing", "supervisor_pid": 11111, "download_pid": None, "shard_pids": [22222], } terminated = {} module._pid_exists = lambda pid: pid in {11111, 22222} module._terminate_pids = lambda pids, grace_seconds=5.0: terminated.setdefault("pids", list(pids)) exit_code, status = module._evaluate( module.parse_args( [ "--state-file", str(tmp_path / "missing-state.json"), "--status-file", str(status_path), "--once", ] ), str(status_path), previous, ) assert exit_code == 6 assert terminated["pids"] == [11111, 22222] assert status["action"] == "terminated" assert status["reason"] == "state_file_missing" def test_watchdog_state_file_missing_scans_for_live_pipeline_candidates(tmp_path): module = _load_module("test_runpod_watchdog_missing_state_scan", "deploy/runpod/watchdog.py") status_path = tmp_path / "watchdog-status.json" terminated = {} module._collect_pipeline_candidate_pids = lambda: [33333, 44444] module._pid_exists = lambda pid: pid in {33333, 44444} module._terminate_pids = lambda pids, grace_seconds=5.0: terminated.setdefault("pids", list(pids)) exit_code, status = module._evaluate( module.parse_args( [ "--state-file", str(tmp_path / "missing-state.json"), "--status-file", str(status_path), "--once", ] ), str(status_path), {}, ) assert exit_code == 6 assert terminated["pids"] == [33333, 44444] assert status["action"] == "terminated" assert status["reason"] == "state_file_missing" def test_watchdog_collects_pipeline_candidate_pids_from_pgrep(monkeypatch): module = _load_module("test_runpod_watchdog_candidate_scan", "deploy/runpod/watchdog.py") def _fake_run(cmd, capture_output, text, check): pattern = cmd[-1] if "run-batch.sh" in pattern: return type("Result", (), {"returncode": 0, "stdout": "111\n222\n"})() if "process-local-videos.py" in pattern: return type("Result", (), {"returncode": 0, "stdout": "222\n333\n"})() return type("Result", (), {"returncode": 1, "stdout": ""})() monkeypatch.setattr(module.subprocess, "run", _fake_run) monkeypatch.setattr(module.os, "getpid", lambda: 333) assert module._collect_pipeline_candidate_pids() == [111, 222] def test_watchdog_pid_exists_treats_zombies_as_dead(monkeypatch): module = _load_module("test_runpod_watchdog_pid_exists", "deploy/runpod/watchdog.py") monkeypatch.setattr(module.Path, "read_text", lambda self, encoding="utf-8": "123 (python) Z 1 2 3") monkeypatch.setattr(module.os, "kill", lambda pid, sig: None) assert module._pid_exists(12345) is False def test_watchdog_terminates_when_all_children_are_dead(tmp_path): module = _load_module("test_runpod_watchdog_all_children_dead", "deploy/runpod/watchdog.py") state_path = tmp_path / "watchdog-state.json" status_path = tmp_path / "watchdog-status.json" state_path.write_text( json.dumps( { "phase": "processing", "heartbeat_epoch": int(time.time()), "supervisor_pid": os.getpid(), "download_pid": 999991, "shard_pids": [999992, 999993], "report_dir": str(tmp_path), "workspace_dir": str(tmp_path), "download_finished": False, "download_exit_code": 0, } ) ) terminated = {} module._terminate_pids = lambda pids, grace_seconds=5.0: terminated.setdefault("pids", list(pids)) exit_code = module.main( [ "--state-file", str(state_path), "--status-file", str(status_path), "--once", ] ) assert exit_code == 7 assert terminated["pids"] == [os.getpid()] status = json.loads(status_path.read_text()) assert status["action"] == "terminated" assert status["reason"] == "all_children_dead" def test_watchdog_detects_export_stall_before_status_file_exists(tmp_path): module = _load_module("test_runpod_watchdog_export_stall", "deploy/runpod/watchdog.py") original_export_status = module.EXPORT_STATUS_FILE module.EXPORT_STATUS_FILE = str(tmp_path / "missing-export-status.json") try: state_path = tmp_path / "watchdog-state.json" status_path = tmp_path / "watchdog-status.json" state_path.write_text( json.dumps( { "phase": "auto-export", "heartbeat_epoch": int(time.time()), "supervisor_pid": os.getpid(), "download_pid": None, "shard_pids": [], "report_dir": str(tmp_path), "workspace_dir": str(tmp_path), "download_finished": True, "download_exit_code": 0, } ) ) args = module.parse_args( [ "--state-file", str(state_path), "--status-file", str(status_path), "--stall-checks", "2", "--once", ] ) terminated = {} module._terminate_pids = lambda pids, grace_seconds=5.0: terminated.setdefault("pids", list(pids)) exit_code, previous = module._evaluate(args, str(status_path), {}) assert exit_code == 0 assert previous["stall_streak"] == 1 exit_code, previous = module._evaluate(args, str(status_path), previous) assert exit_code == 5 assert terminated["pids"] == [os.getpid()] assert previous["action"] == "terminated" assert previous["reason"] == "export_stall" finally: module.EXPORT_STATUS_FILE = original_export_status def test_watchdog_detects_export_stall_when_status_and_signature_freeze(tmp_path): module = _load_module("test_runpod_watchdog_export_frozen", "deploy/runpod/watchdog.py") state_path = tmp_path / "watchdog-state.json" status_path = tmp_path / "watchdog-status.json" state_path.write_text( json.dumps( { "phase": "auto-export", "heartbeat_epoch": int(time.time()), "supervisor_pid": os.getpid(), "download_pid": None, "shard_pids": [], "report_dir": str(tmp_path), "workspace_dir": str(tmp_path), "download_finished": True, "download_exit_code": 0, } ) ) module._collect_export_status = lambda: {"stage": "tar", "_mtime_ns": 1, "_size": 10} module._collect_export_signature = lambda: ["tarball.tgz", 123, 456] terminated = {} module._terminate_pids = lambda pids, grace_seconds=5.0: terminated.setdefault("pids", list(pids)) args = module.parse_args( [ "--state-file", str(state_path), "--status-file", str(status_path), "--stall-checks", "2", "--once", ] ) exit_code, previous = module._evaluate(args, str(status_path), {}) assert exit_code == 0 assert previous["stall_streak"] == 0 exit_code, previous = module._evaluate(args, str(status_path), previous) assert exit_code == 0 assert previous["stall_streak"] == 1 exit_code, previous = module._evaluate(args, str(status_path), previous) assert exit_code == 5 assert terminated["pids"] == [os.getpid()] assert previous["action"] == "terminated" assert previous["reason"] == "export_stall" def test_restore_runtime_dbs_prefers_newer_valid_backup(tmp_path): module = _load_module("test_restore_runtime_dbs", "deploy/runpod/restore_runtime_dbs.py") db_utils = _load_module("test_restore_runtime_dbs_utils", "backend/db_utils.py") runtime_dir = tmp_path / "runtime" backup_dir = tmp_path / "backup" runtime_dir.mkdir() backup_dir.mkdir() runtime_db = runtime_dir / "images_database.db" backup_db = backup_dir / "images_database.db" runtime_conn = db_utils.connect_vec_db(str(runtime_db)) runtime_conn.execute("CREATE TABLE sample(value TEXT)") runtime_conn.execute("INSERT INTO sample(value) VALUES ('runtime')") runtime_conn.commit() runtime_conn.close() time.sleep(0.01) backup_conn = db_utils.connect_vec_db(str(backup_db)) backup_conn.execute("CREATE TABLE sample(value TEXT)") backup_conn.execute("INSERT INTO sample(value) VALUES ('backup')") backup_conn.commit() backup_conn.close() exit_code = module.main( [ "--runtime-dir", str(runtime_dir), "--backup-dir", str(backup_dir), ] ) assert exit_code == 0 check_conn = db_utils.connect_vec_db(str(runtime_db)) value = check_conn.execute("SELECT value FROM sample").fetchone()[0] check_conn.close() assert value == "backup" def test_restore_runtime_dbs_keeps_newer_valid_runtime(tmp_path): module = _load_module("test_restore_runtime_dbs_keep", "deploy/runpod/restore_runtime_dbs.py") db_utils = _load_module("test_restore_runtime_dbs_keep_utils", "backend/db_utils.py") runtime_dir = tmp_path / "runtime" backup_dir = tmp_path / "backup" runtime_dir.mkdir() backup_dir.mkdir() runtime_db = runtime_dir / "images_database.db" backup_db = backup_dir / "images_database.db" backup_conn = db_utils.connect_vec_db(str(backup_db)) backup_conn.execute("CREATE TABLE sample(value TEXT)") backup_conn.execute("INSERT INTO sample(value) VALUES ('backup')") backup_conn.commit() backup_conn.close() time.sleep(0.01) runtime_conn = db_utils.connect_vec_db(str(runtime_db)) runtime_conn.execute("CREATE TABLE sample(value TEXT)") runtime_conn.execute("INSERT INTO sample(value) VALUES ('runtime')") runtime_conn.commit() runtime_conn.close() exit_code = module.main( [ "--runtime-dir", str(runtime_dir), "--backup-dir", str(backup_dir), ] ) assert exit_code == 0 check_conn = db_utils.connect_vec_db(str(runtime_db)) value = check_conn.execute("SELECT value FROM sample").fetchone()[0] check_conn.close() assert value == "runtime" def test_backup_runtime_dbs_writes_manifest(tmp_path): module = _load_module("test_backup_runtime_dbs_manifest", "deploy/runpod/backup_runtime_dbs.py") db_utils = _load_module("test_backup_runtime_dbs_manifest_utils", "backend/db_utils.py") source_dir = tmp_path / "runtime" target_dir = tmp_path / "backup" source_dir.mkdir() target_dir.mkdir() source_db = source_dir / "images_database.db" conn = db_utils.connect_vec_db(str(source_db)) conn.execute("CREATE TABLE sample(value TEXT)") conn.execute("INSERT INTO sample(value) VALUES ('runtime')") conn.commit() conn.close() exit_code = module.main( [ "--source-dir", str(source_dir), "--target-dir", str(target_dir), ] ) assert exit_code == 0 manifest = json.loads((target_dir / "backup-manifest.json").read_text()) entry = manifest["databases"]["images_database.db"] assert entry["quick_check"] == "ok" assert entry["size_bytes"] > 0 assert entry["sha256"] def test_restore_runtime_dbs_rejects_manifest_checksum_mismatch(tmp_path): backup_module = _load_module("test_backup_runtime_dbs_manifest_restore", "deploy/runpod/backup_runtime_dbs.py") restore_module = _load_module("test_restore_runtime_dbs_manifest_restore", "deploy/runpod/restore_runtime_dbs.py") db_utils = _load_module("test_restore_runtime_dbs_manifest_restore_utils", "backend/db_utils.py") runtime_dir = tmp_path / "runtime" backup_dir = tmp_path / "backup" runtime_dir.mkdir() backup_dir.mkdir() source_db = runtime_dir / "images_database.db" conn = db_utils.connect_vec_db(str(source_db)) conn.execute("CREATE TABLE sample(value TEXT)") conn.execute("INSERT INTO sample(value) VALUES ('runtime')") conn.commit() conn.close() backup_module.main( [ "--source-dir", str(runtime_dir), "--target-dir", str(backup_dir), ] ) backup_db = backup_dir / "images_database.db" backup_db.write_bytes(backup_db.read_bytes() + b"corrupt") source_db.unlink() exit_code = restore_module.main( [ "--runtime-dir", str(runtime_dir), "--backup-dir", str(backup_dir), ] ) assert exit_code == 1 assert not source_db.exists() # --------------------------------------------------------------------------- # _chat_with_retry — backoff on transient API errors (429/5xx/timeout) # --------------------------------------------------------------------------- class _FakeStatusError(Exception): """Mimics an openai SDK error carrying a .status_code attribute.""" def __init__(self, status_code, message=""): super().__init__(message or f"status {status_code}") self.status_code = status_code class _FakeClient: """A client whose chat.completions.create raises a queued sequence of exceptions, then returns a sentinel once the queue is exhausted.""" def __init__(self, errors): self._errors = list(errors) self.calls = 0 outer = self class _Completions: def create(self, **kwargs): outer.calls += 1 if outer._errors: raise outer._errors.pop(0) return "OK_RESPONSE" class _Chat: completions = _Completions() self.chat = _Chat() def _load_runpod_processor(monkeypatch): module = _load_module("test_runpod_scene_processor", "scripts/runpod_scene_processor.py") # Make backoff instant so the test doesn't actually sleep. monkeypatch.setattr(module.time, "sleep", lambda *_a, **_k: None) return module def test_chat_with_retry_recovers_after_429(monkeypatch): module = _load_runpod_processor(monkeypatch) client = _FakeClient([_FakeStatusError(429), _FakeStatusError(429)]) result = module._chat_with_retry(client, model="m", messages=[]) assert result == "OK_RESPONSE" assert client.calls == 3 # two failures + one success def test_chat_with_retry_reraises_fatal_400_immediately(monkeypatch): module = _load_runpod_processor(monkeypatch) client = _FakeClient([_FakeStatusError(400, "bad request")]) try: module._chat_with_retry(client, model="m", messages=[]) raise AssertionError("expected the 400 to be re-raised") except _FakeStatusError as exc: assert exc.status_code == 400 assert client.calls == 1 # no retries on a non-transient error def test_chat_with_retry_gives_up_after_max_retries(monkeypatch): module = _load_runpod_processor(monkeypatch) # Always 429 — should exhaust max_retries (8) and then raise. client = _FakeClient([_FakeStatusError(429) for _ in range(20)]) try: module._chat_with_retry(client, model="m", messages=[], max_retries=8) raise AssertionError("expected to give up and raise after max retries") except _FakeStatusError as exc: assert exc.status_code == 429 assert client.calls == 8 # exactly max_retries attempts, no infinite loop def test_chat_with_retry_retries_on_timeout_message(monkeypatch): module = _load_runpod_processor(monkeypatch) # No status_code, but message indicates a timeout → should retry. client = _FakeClient([TimeoutError("request timeout")]) result = module._chat_with_retry(client, model="m", messages=[]) assert result == "OK_RESPONSE" assert client.calls == 2