"""Unit tests for ``generate_with_batch_management`` wrapper behavior.""" import inspect import unittest from unittest.mock import patch from _batch_management_test_support import build_progress_result from _batch_management_test_support import load_batch_management_module def _build_call_kwargs(module): """Build complete kwargs for ``generate_with_batch_management``.""" kwargs = {} for name in list(inspect.signature(module.generate_with_batch_management).parameters)[2:]: if name == "progress": continue if name == "batch_size_input": kwargs[name] = 2 elif name in ("allow_lm_batch", "auto_lrc", "autogen_checkbox", "auto_score"): kwargs[name] = False elif name == "current_batch_index": kwargs[name] = 0 elif name == "total_batches": kwargs[name] = 0 elif name in ("batch_queue", "generation_params_state"): kwargs[name] = {} elif name == "complete_track_classes": kwargs[name] = [] else: kwargs[name] = None return kwargs class BatchManagementWrapperTests(unittest.TestCase): """Tests for streaming and final wrapper output mapping.""" def test_non_windows_streams_partial_and_final_outputs(self): """Non-Windows path should emit partial UI updates plus final state.""" module, state = load_batch_management_module(is_windows=False) def _gen(*_args, **_kwargs): """Yield one standard progress result for wrapper streaming.""" yield build_progress_result(length=48) kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): outputs = list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual(len(outputs), 2) self.assertEqual(len(outputs[0]), 55) self.assertEqual(len(outputs[1]), 55) self.assertEqual(outputs[0][0]["playback_position"], 0) self.assertEqual(outputs[1][0]["playback_position"], 0) self.assertEqual(len(state["store_calls"]), 1) def test_windows_emits_only_final_output(self): """Windows path should skip intermediate yields and emit final state only.""" module, _state = load_batch_management_module(is_windows=True) def _gen(*_args, **_kwargs): """Yield one standard progress result for Windows final-output path.""" yield build_progress_result(length=48) kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): outputs = list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual(len(outputs), 1) self.assertEqual(len(outputs[0]), 55) self.assertEqual(outputs[0][0]["playback_position"], 0) def test_all_audio_paths_none_skips_batch_storage(self): """When inner result has no audio paths, wrapper should not store a batch.""" module, state = load_batch_management_module(is_windows=False) def _gen(*_args, **_kwargs): """Yield a result with no audio paths to trigger early return.""" yield build_progress_result(length=48, all_audio_paths=None) kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): outputs = list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual(len(outputs), 2) self.assertEqual(outputs[1][8], None) self.assertEqual(len(state["store_calls"]), 0) def test_allow_lm_batch_stores_multiple_codes(self): """Batch mode should store a list of generated codes up to batch size.""" module, state = load_batch_management_module(is_windows=False) def _gen(*_args, **_kwargs): """Yield a result carrying a list of generated codes.""" result = list(build_progress_result(length=48)) result[47] = [f"code-{idx}" for idx in range(8)] yield tuple(result) kwargs = _build_call_kwargs(module) kwargs["allow_lm_batch"] = True kwargs["batch_size_input"] = 3 with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual(state["store_calls"][0]["codes"], ["code-0", "code-1", "code-2"]) def test_no_fsq_forwards_to_generation_and_saved_params(self): """Wrapper should pass and persist the Remix no_fsq checkbox.""" module, state = load_batch_management_module(is_windows=False) seen = {} def _gen(*args, **_kwargs): """Capture positional generation args and yield a standard result.""" seen["args"] = args yield build_progress_result(length=48) kwargs = _build_call_kwargs(module) kwargs["task_type"] = "cover" kwargs["no_fsq"] = True with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): list(module.generate_with_batch_management(None, None, **kwargs)) self.assertTrue(seen["args"][23]) saved_params = state["store_calls"][0]["generation_params"] self.assertTrue(saved_params["no_fsq"]) self.assertEqual(saved_params["task_type"], "cover") def test_auto_lrc_copies_lrc_fields_to_batch_queue(self): """Auto-LRC mode should copy LRC/subtitle payloads into stored queue entry.""" module, _state = load_batch_management_module(is_windows=False) lrcs = [f"lrc-{idx}" for idx in range(8)] subtitles = [f"sub-{idx}" for idx in range(8)] def _gen(*_args, **_kwargs): """Yield a result with explicit LRC/subtitle payload.""" result = list(build_progress_result(length=48)) result[46] = {"lrcs": lrcs, "subtitles": subtitles} yield tuple(result) kwargs = _build_call_kwargs(module) kwargs["auto_lrc"] = True outputs = [] with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): outputs = list(module.generate_with_batch_management(None, None, **kwargs)) final_batch_queue = outputs[-1][48] self.assertEqual(final_batch_queue[0]["lrcs"], lrcs) self.assertEqual(final_batch_queue[0]["subtitles"], subtitles) def test_auto_lrc_sets_lrc_display_in_final_yield(self): """Final yield should carry gr.update(value=lrc) at positions 36-43.""" module, _state = load_batch_management_module(is_windows=False) lrcs = [f"[00:01.00]Line {idx}" for idx in range(8)] subtitles = [f"sub-{idx}" for idx in range(8)] def _gen(*_args, **_kwargs): """Yield a result with LRC data in extra_outputs.""" result = list(build_progress_result(length=48)) result[46] = {"lrcs": lrcs, "subtitles": subtitles} yield tuple(result) kwargs = _build_call_kwargs(module) kwargs["auto_lrc"] = True with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): outputs = list(module.generate_with_batch_management(None, None, **kwargs)) final_yield = outputs[-1] for i in range(8): lrc_val = final_yield[36 + i] self.assertIsInstance(lrc_val, dict, f"LRC position {36 + i} should be a gr.update dict") self.assertEqual( lrc_val.get("value"), lrcs[i], f"LRC position {36 + i} should contain the LRC text", ) def test_auto_lrc_disabled_preserves_passthrough_values(self): """When auto_lrc is off, LRC positions pass through from inner generator.""" module, _state = load_batch_management_module(is_windows=False) def _gen(*_args, **_kwargs): """Yield a standard result without auto_lrc.""" yield build_progress_result(length=48) kwargs = _build_call_kwargs(module) kwargs["auto_lrc"] = False with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): outputs = list(module.generate_with_batch_management(None, None, **kwargs)) final_yield = outputs[-1] for i in range(8): lrc_val = final_yield[36 + i] self.assertIsNone(lrc_val, f"LRC position {36 + i} should be None when auto_lrc is off") def test_empty_inner_generator_returns_skip_tuple_and_warning(self): """Empty inner generator should fail gracefully without indexing None.""" module, state = load_batch_management_module(is_windows=False) def _gen(*_args, **_kwargs): """Yield nothing to simulate a defensive empty-generator edge case.""" if False: yield None kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): outputs = list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual(len(outputs), 1) self.assertEqual(len(outputs[0]), 55) self.assertTrue(all(item.get("kind") == "skip" for item in outputs[0])) self.assertEqual(len(state["store_calls"]), 0) self.assertTrue(state["warning_messages"]) self.assertIn("messages.batch_failed", state["warning_messages"][0]) # ------------------------------------------------------------------ # Score persistence regression tests (foreground batch fix) # ------------------------------------------------------------------ def test_foreground_scores_passed_to_store_batch_in_queue(self): """Foreground generation must extract and pass scores to batch storage.""" module, state = load_batch_management_module(is_windows=False) def _gen(*_args, **_kwargs): """Yield a result with score values at indices 12-19.""" result = list(build_progress_result(length=48)) for i in range(8): result[12 + i] = f"8.{i}" yield tuple(result) kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual(len(state["store_calls"]), 1) scores = state["store_calls"][0]["scores"] self.assertEqual(len(scores), 8) self.assertEqual(scores[0], "8.0") self.assertEqual(scores[7], "8.7") def test_foreground_scores_default_empty_when_absent(self): """When result tuple lacks score indices, scores should be empty strings.""" module, state = load_batch_management_module(is_windows=False) def _gen(*_args, **_kwargs): """Yield a short result with no score data.""" yield build_progress_result(length=48) kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual(len(state["store_calls"]), 1) scores = state["store_calls"][0]["scores"] self.assertEqual(len(scores), 8) self.assertTrue(all(s == "" for s in scores), "Absent scores should default to empty strings") # ------------------------------------------------------------------ # MPS cache-clearing regression tests (macOS audio-mute fix) # ------------------------------------------------------------------ def test_mps_cache_cleared_before_and_after_generation_on_mac(self): """On MPS, empty_cache must be called both before and after generation.""" module, state = load_batch_management_module(is_windows=False, mps_available=True) def _gen(*_args, **_kwargs): """Yield one result for MPS cache-clearing path.""" yield build_progress_result(length=48) kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): list(module.generate_with_batch_management(None, None, **kwargs)) self.assertGreaterEqual( state["mps_empty_cache_calls"], 2, "torch.mps.empty_cache() must be called before and after generation " "on macOS to prevent system audio mute", ) def test_mps_cache_not_called_when_mps_unavailable(self): """MPS cache clear must not be called when MPS is absent (non-Mac hosts).""" module, state = load_batch_management_module(is_windows=False, mps_available=False) def _gen(*_args, **_kwargs): """Yield one result for non-MPS path.""" yield build_progress_result(length=48) kwargs = _build_call_kwargs(module) with patch.dict(module.generate_with_batch_management.__globals__, {"generate_with_progress": _gen}): list(module.generate_with_batch_management(None, None, **kwargs)) self.assertEqual( state["mps_empty_cache_calls"], 0, "torch.mps.empty_cache() must not be called when MPS is unavailable", ) if __name__ == "__main__": unittest.main()