diff --git a/abogen/domain/audio_sink.py b/abogen/domain/audio_sink.py new file mode 100644 index 0000000..0fe9bd7 --- /dev/null +++ b/abogen/domain/audio_sink.py @@ -0,0 +1,131 @@ +"""Audio sink abstraction for unified audio output. + +Provides a context-manager-based abstraction for writing audio data +to various output formats (WAV, FLAC via soundfile; compressed via ffmpeg). + +Usage: + with open_audio_sink(path, "wav") as sink: + sink.write(audio_data) +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Optional + +import numpy as np + +from abogen.domain.audio_buffer import SAMPLE_RATE +from abogen.domain.audio_helpers import build_ffmpeg_command + + +@dataclass(frozen=True) +class AudioSink: + """Represents an open audio output target.""" + + write: Callable[[np.ndarray], None] + close: Callable[[], None] + + def __enter__(self) -> AudioSink: + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.close() + + +def _ensure_ffmpeg() -> None: + """Ensure static ffmpeg binaries are on PATH.""" + import static_ffmpeg # type: ignore + + ffmpeg_cache_root = _get_ffmpeg_cache_root() + platform_cache = os.path.join(ffmpeg_cache_root, sys.platform) + os.makedirs(platform_cache, exist_ok=True) + try: + import static_ffmpeg.run as static_ffmpeg_run # type: ignore + + static_ffmpeg_run.LOCK_FILE = os.path.join(ffmpeg_cache_root, "lock.file") + except Exception: + pass + static_ffmpeg.add_paths(weak=True, download_dir=platform_cache) + + +def _get_ffmpeg_cache_root() -> str: + from abogen.infrastructure.cache import get_internal_cache_path + + return get_internal_cache_path("ffmpeg") + + +def open_audio_sink( + path: Path, + fmt: str, + *, + metadata: Optional[dict[str, str]] = None, + cancel_check: Optional[Callable[[], bool]] = None, + extra_ffmpeg_args: Optional[list[str]] = None, + ffmpeg_cmd: Optional[list[str]] = None, +) -> AudioSink: + """Open an audio output sink for writing raw float32 PCM samples. + + Args: + path: Output file path. + fmt: Output format ("wav", "flac", "mp3", "opus", "m4b"). + metadata: Optional metadata dict (ignored when ffmpeg_cmd is provided). + cancel_check: Optional callable; if it returns True, writes are silently skipped. + extra_ffmpeg_args: Optional extra args inserted after ffmpeg header (ignored when ffmpeg_cmd is provided). + ffmpeg_cmd: Optional pre-built ffmpeg command list (for m4b with cover art etc.). + + Returns: + AudioSink with write() and close() methods. + """ + fmt = fmt.lower() + + if fmt in {"wav", "flac"}: + import soundfile as sf + + soundfile_obj = sf.SoundFile( + path, + mode="w", + samplerate=SAMPLE_RATE, + channels=1, + format=fmt.upper(), + ) + + def _write_wav(data: np.ndarray) -> None: + if cancel_check and cancel_check(): + return + soundfile_obj.write(data) + + def _close_wav() -> None: + soundfile_obj.close() + + return AudioSink(write=_write_wav, close=_close_wav) + + # Compressed formats: pipe through ffmpeg + _ensure_ffmpeg() + + if ffmpeg_cmd is not None: + cmd = list(ffmpeg_cmd) + else: + cmd = build_ffmpeg_command(path, fmt, metadata=metadata) + if extra_ffmpeg_args: + cmd[2:2] = extra_ffmpeg_args + + process = subprocess.Popen( + cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL + ) + + def _write_compressed(data: np.ndarray) -> None: + if (cancel_check and cancel_check()) or process.stdin is None or process.stdin.closed: + return + process.stdin.write(data.tobytes()) + + def _close_compressed() -> None: + if process.stdin and not process.stdin.closed: + process.stdin.close() + process.wait() + + return AudioSink(write=_write_compressed, close=_close_compressed) diff --git a/abogen/pyqt/conversion.py b/abogen/pyqt/conversion.py index 651b5f5..3a86005 100644 --- a/abogen/pyqt/conversion.py +++ b/abogen/pyqt/conversion.py @@ -5,6 +5,7 @@ import hashlib # For generating unique cache filenames from platformdirs import user_desktop_dir from PyQt6.QtCore import QThread, pyqtSignal, Qt, QTimer from PyQt6.QtWidgets import QCheckBox, QVBoxLayout, QDialog, QLabel, QDialogButtonBox +from contextlib import ExitStack import numpy as np import soundfile as sf from abogen.utils import ( @@ -29,6 +30,7 @@ from abogen.domain.output_paths import ( sanitize_output_stem, ) from abogen.domain.audio_helpers import build_ffmpeg_command, to_float32 +from abogen.domain.audio_sink import AudioSink, open_audio_sink from abogen.domain.audio_buffer import ( create_silence, mix_audio, @@ -361,6 +363,7 @@ class ConversionThread(QThread): ) try: hf_tracker.set_log_callback(lambda msg: self.log_updated.emit(msg)) + sink_stack = ExitStack() # Show configuration self.log_updated.emit("Configuration:") @@ -735,57 +738,12 @@ class ConversionThread(QThread): # SRT numbering fix: use a global counter merged_srt_index = 1 # SRT numbering for merged file # Prepare output file/ffmpeg process for merged output - if self.output_format in ["wav", "mp3", "flac"]: - merged_out_file = sf.SoundFile( + merged_sink = sink_stack.enter_context( + self._open_merged_sink( merged_out_path, - "w", - samplerate=24000, - channels=1, - format=self.output_format, + cancel_check=lambda: self.cancel_requested, ) - ffmpeg_proc = None - elif self.output_format in ("m4b", "opus"): - # Real-time generation using FFmpeg pipe - static_ffmpeg.add_paths() - merged_out_file = None - ffmpeg_proc = None - metadata_options, cover_path = ( - self._extract_and_add_metadata_tags_to_ffmpeg_cmd() - if self.output_format == "m4b" - else ([], None) - ) - cmd = build_ffmpeg_command( - Path(merged_out_path), - self.output_format, - ) - # Insert thread queue size after ffmpeg header - cmd.insert(2, "-thread_queue_size") - cmd.insert(3, "32768") - if self.output_format == "m4b" and cover_path and os.path.exists(cover_path): - # Insert cover image input before the output path - output_path = cmd.pop() - cmd.extend([ - "-i", cover_path, - "-map", "0:a", - "-map", "1", - "-c:v", "copy", - "-disposition:v", "attached_pic", - ]) - cmd.extend(metadata_options) - cmd.append(output_path) - elif self.output_format == "m4b": - output_path = cmd.pop() - cmd.extend(metadata_options) - cmd.append(output_path) - ffmpeg_proc = create_process(cmd, stdin=subprocess.PIPE, text=False) - else: - self.log_updated.emit( - (f"Unsupported output format: {self.output_format}", "red") - ) - self.conversion_finished.emit( - ("Audio generation failed.", "red"), None - ) - return + ) # Open merged subtitle file for incremental writing if needed merged_subtitle_file = None if self.subtitle_mode != "Disabled": @@ -849,9 +807,8 @@ class ConversionThread(QThread): merged_subtitle_path = None merged_subtitle_file = None else: - # If not merging, set merged_out_file and related variables to None - merged_out_file = None - ffmpeg_proc = None + # If not merging, set merged_sink and related variables to None + merged_sink = None merged_out_path = None subtitle_entries = [] current_time = 0.0 @@ -868,8 +825,7 @@ class ConversionThread(QThread): # Instead of processing the whole text, process by chapter for chapter_idx, (chapter_name, voice_segments) in enumerate(chapters, 1): chapter_out_path = None - chapter_out_file = None - chapter_ffmpeg_proc = None + chapter_sink = None chapter_subtitle_file = None chapter_subtitle_path = None if total_chapters > 1: @@ -904,37 +860,12 @@ class ConversionThread(QThread): chapters_out_dir, f"{chapter_filename}.{separate_chapters_format}", ) - if separate_chapters_format in ["wav", "mp3", "flac"]: - chapter_out_file = sf.SoundFile( - chapter_out_path, - "w", - samplerate=24000, - channels=1, - format=separate_chapters_format, - ) - chapter_ffmpeg_proc = None - elif separate_chapters_format == "opus": - static_ffmpeg.add_paths() - cmd = [ - "ffmpeg", - "-y", - "-thread_queue_size", - "32768", - "-f", - "f32le", - "-ar", - "24000", - "-ac", - "1", - "-i", - "pipe:0", - ] - cmd.extend(["-c:a", "libopus", "-b:a", "24000"]) - cmd.append(chapter_out_path) - chapter_ffmpeg_proc = create_process( - cmd, stdin=subprocess.PIPE, text=False - ) - chapter_out_file = None + if separate_chapters_format in ("wav", "flac", "mp3", "opus"): + chapter_sink = sink_stack.enter_context(open_audio_sink( + Path(chapter_out_path), + separate_chapters_format, + cancel_check=lambda: self.cancel_requested, + )) else: self.log_updated.emit( ( @@ -1129,10 +1060,7 @@ class ConversionThread(QThread): # Print the result for debugging # print(f"Result: {result}") if self.cancel_requested: - if chapter_out_file: - chapter_out_file.close() - if merged_out_file: - merged_out_file.close() + sink_stack.close() self.conversion_finished.emit("Cancelled", None) return current_segment += 1 @@ -1147,14 +1075,10 @@ class ConversionThread(QThread): chunk_start = current_time audio_np = to_float32(result.audio) # Write audio directly to merged file ONLY if merging - if merge_chapters_at_end and merged_out_file: - merged_out_file.write(audio_np) - elif merge_chapters_at_end and ffmpeg_proc: - ffmpeg_proc.stdin.write(audio_np.tobytes()) - if chapter_out_file: - chapter_out_file.write(audio_np) - elif chapter_ffmpeg_proc: - chapter_ffmpeg_proc.stdin.write(audio_np.tobytes()) + if merge_chapters_at_end and merged_sink: + merged_sink.write(audio_np) + if chapter_sink: + chapter_sink.write(audio_np) # Subtitle logic if self.subtitle_mode != "Disabled": tokens_list = getattr(result, "tokens", []) @@ -1187,7 +1111,7 @@ class ConversionThread(QThread): "whitespace": tok.whitespace, } ) - if chapter_out_file or chapter_ffmpeg_proc: + if chapter_sink: chapter_tokens_with_timestamps.append( { "start": chapter_current_time @@ -1234,8 +1158,8 @@ class ConversionThread(QThread): f"{merged_srt_index}\n{_format_timestamp(start)} --> {_format_timestamp(end)}\n{text}\n\n" ) merged_srt_index += 1 - # Per-chapter subtitle processing for both file and ffmpeg_proc - if chapter_out_file or chapter_ffmpeg_proc: + # Per-chapter subtitle processing for both file and sink + if chapter_sink: new_chapter_entries = [] self._process_subtitle_tokens( chapter_tokens_with_timestamps, @@ -1270,10 +1194,10 @@ class ConversionThread(QThread): chapter_srt_index += 1 if merge_chapters_at_end: current_time += chunk_dur - if chapter_out_file or chapter_ffmpeg_proc: + if chapter_sink: chapter_current_time += chunk_dur else: - if chapter_out_file or chapter_ffmpeg_proc: + if chapter_sink: chapter_current_time += chunk_dur # Calculate percentage based on characters processed percent = min( @@ -1296,29 +1220,22 @@ class ConversionThread(QThread): # Add silence between chapters for merged output (except after the last chapter) if merge_chapters_at_end and chapter_idx < total_chapters: silence_audio = create_silence(self.silence_duration) - silence_bytes = silence_audio.tobytes() - if merged_out_file: - merged_out_file.write(silence_audio) - elif ffmpeg_proc: - ffmpeg_proc.stdin.write(silence_bytes) + if merged_sink: + merged_sink.write(silence_audio) + # Update timing for the silence current_time += self.silence_duration - if chapter_out_file or chapter_ffmpeg_proc: + if chapter_sink: chapter_current_time += self.silence_duration # Set chapter end time after processing if merge_chapters_at_end: chapter_time["end"] = current_time # Finalize chapter file for ffmpeg formats - if chapter_out_file or chapter_ffmpeg_proc: + if chapter_sink: self.log_updated.emit(("\nProcessing chapter audio...", "grey")) - if chapter_ffmpeg_proc: - chapter_ffmpeg_proc.stdin.close() - chapter_ffmpeg_proc.wait() - if chapter_out_file: - chapter_out_file.close() # Close chapter subtitle file if open if chapter_subtitle_file: chapter_subtitle_file.close() @@ -1344,11 +1261,8 @@ class ConversionThread(QThread): # Finalize merged output file ONLY if merging if merge_chapters_at_end: self.log_updated.emit(("\nFinalizing audio. Please wait...", "grey")) - if self.output_format in ["wav", "mp3", "flac"]: - merged_out_file.close() - elif self.output_format == "m4b": - ffmpeg_proc.stdin.close() - ffmpeg_proc.wait() + merged_sink.close() + if self.output_format == "m4b": # Add chapters via fast post-processing if total_chapters > 1: chapters_info_path = f"{base_filepath_no_ext}_chapters.txt" @@ -1412,8 +1326,7 @@ class ConversionThread(QThread): os.replace(tmp_path, orig_path) os.remove(chapters_info_path) elif self.output_format in ["opus"]: - ffmpeg_proc.stdin.close() - ffmpeg_proc.wait() + merged_sink.close() self.progress_updated.emit(100, "00:00:00") # Close merged subtitle file if open if merged_subtitle_file: @@ -1442,23 +1355,15 @@ class ConversionThread(QThread): chapters_dir, ) except Exception as e: - # Cleanup ffmpeg subprocesses on error + # Cleanup all sinks on error try: - if "ffmpeg_proc" in locals() and ffmpeg_proc: - ffmpeg_proc.stdin.close() - ffmpeg_proc.terminate() - ffmpeg_proc.wait() - except Exception: - pass - try: - if "chapter_ffmpeg_proc" in locals() and chapter_ffmpeg_proc: - chapter_ffmpeg_proc.stdin.close() - chapter_ffmpeg_proc.terminate() - chapter_ffmpeg_proc.wait() + sink_stack.close() except Exception: pass self.log_updated.emit((f"Error occurred: {str(e)}", "red")) self.conversion_finished.emit(("Audio generation failed.", "red"), None) + finally: + sink_stack.close() def _process_subtitle_file(self, tts, base_path, is_timestamp_text=False): """Process subtitle files with precise timing and generate output subtitles.""" @@ -1527,39 +1432,10 @@ class ConversionThread(QThread): rate = 24000 # Setup audio output - merged_out_file, ffmpeg_proc = None, None - if self.output_format in ["wav", "mp3", "flac"]: - merged_out_file = sf.SoundFile( - merged_out_path, - "w", - samplerate=rate, - channels=1, - format=self.output_format, - ) - else: - static_ffmpeg.add_paths() - cmd = build_ffmpeg_command( - Path(merged_out_path), - self.output_format, - ) - cmd.insert(2, "-thread_queue_size") - cmd.insert(3, "32768") - if self.output_format == "m4b": - metadata_options, cover_path = ( - self._extract_and_add_metadata_tags_to_ffmpeg_cmd() - ) - if cover_path and os.path.exists(cover_path): - output_path = cmd.pop() - cmd.extend([ - "-i", cover_path, - "-map", "0:a", - "-map", "1", - "-c:v", "copy", - "-disposition:v", "attached_pic", - ]) - cmd.append(output_path) - cmd.extend(metadata_options) - ffmpeg_proc = create_process(cmd, stdin=subprocess.PIPE, text=False) + merged_sink = self._open_merged_sink( + merged_out_path, + cancel_check=lambda: self.cancel_requested, + ).__enter__() # Always generate subtitles for subtitle input files subtitle_file, subtitle_path = None, None @@ -1851,13 +1727,9 @@ class ConversionThread(QThread): # Write the complete audio buffer self.log_updated.emit(("\nFinalizing audio. Please wait...", "grey")) - if merged_out_file: - merged_out_file.write(audio_buffer) - merged_out_file.close() - elif ffmpeg_proc: - ffmpeg_proc.stdin.write(to_float32(audio_buffer).tobytes()) - ffmpeg_proc.stdin.close() - ffmpeg_proc.wait() + if merged_sink: + merged_sink.write(to_float32(audio_buffer)) + merged_sink.close() if subtitle_file: subtitle_file.close() @@ -1870,10 +1742,8 @@ class ConversionThread(QThread): except Exception as e: try: - if "ffmpeg_proc" in locals() and ffmpeg_proc: - ffmpeg_proc.stdin.close() - ffmpeg_proc.terminate() - ffmpeg_proc.wait() + if "merged_sink" in locals() and merged_sink: + merged_sink.close() if "subtitle_file" in locals() and subtitle_file: subtitle_file.close() except: @@ -1893,6 +1763,52 @@ class ConversionThread(QThread): self._timestamp_response = treat_as_subtitle self._timestamp_response_event.set() + @contextmanager + def _open_merged_sink(self, path, cancel_check=None): + """Open audio sink for the merged output, handling m4b cover art.""" + if self.output_format in ("wav", "flac", "mp3", "opus"): + with open_audio_sink( + Path(path), + self.output_format, + cancel_check=cancel_check, + ) as sink: + yield sink + elif self.output_format == "m4b": + static_ffmpeg.add_paths() + metadata_options, cover_path = ( + self._extract_and_add_metadata_tags_to_ffmpeg_cmd() + ) + cmd = build_ffmpeg_command( + Path(path), + self.output_format, + ) + cmd.insert(2, "-thread_queue_size") + cmd.insert(3, "32768") + if cover_path and os.path.exists(cover_path): + output_path = cmd.pop() + cmd.extend([ + "-i", cover_path, + "-map", "0:a", + "-map", "1", + "-c:v", "copy", + "-disposition:v", "attached_pic", + ]) + cmd.extend(metadata_options) + cmd.append(output_path) + else: + output_path = cmd.pop() + cmd.extend(metadata_options) + cmd.append(output_path) + with open_audio_sink( + Path(path), + self.output_format, + ffmpeg_cmd=cmd, + cancel_check=cancel_check, + ) as sink: + yield sink + else: + raise ValueError(f"Unsupported output format: {self.output_format}") + def _extract_and_add_metadata_tags_to_ffmpeg_cmd(self): """Extract metadata tags from text content and add them to ffmpeg command""" # Read text for metadata extraction @@ -1960,21 +1876,6 @@ class ConversionThread(QThread): self.process.terminate() except Exception: pass - # Terminate ffmpeg subprocesses if running - try: - if hasattr(self, "ffmpeg_proc") and self.ffmpeg_proc: - self.ffmpeg_proc.stdin.close() - self.ffmpeg_proc.terminate() - self.ffmpeg_proc.wait() - except Exception: - pass - try: - if hasattr(self, "chapter_ffmpeg_proc") and self.chapter_ffmpeg_proc: - self.chapter_ffmpeg_proc.stdin.close() - self.chapter_ffmpeg_proc.terminate() - self.chapter_ffmpeg_proc.wait() - except Exception: - pass class VoicePreviewThread(QThread): diff --git a/abogen/webui/conversion_runner.py b/abogen/webui/conversion_runner.py index 5fd1afc..4e0a728 100644 --- a/abogen/webui/conversion_runner.py +++ b/abogen/webui/conversion_runner.py @@ -2,8 +2,6 @@ from __future__ import annotations import json import os -import subprocess -import sys import time import traceback import gc @@ -14,8 +12,6 @@ from pathlib import Path from typing import Any, Callable, Dict, List, Mapping, Optional import numpy as np -import soundfile as sf -import static_ffmpeg from abogen.tts_plugin.utils import is_plugin_registered from abogen.infrastructure.exporters import ExportService @@ -124,6 +120,7 @@ from abogen.domain.audio_buffer import ( normalize_audio as _normalize_audio, SAMPLE_RATE, ) +from abogen.domain.audio_sink import AudioSink, open_audio_sink from .service import Job, JobStatus @@ -149,11 +146,6 @@ class _JobCancelled(Exception): """Raised internally to abort a conversion when the client cancels.""" -@dataclass -class AudioSink: - write: Callable[[np.ndarray], None] - - _APOSTROPHE_CONFIG = ApostropheConfig() @@ -402,7 +394,14 @@ def run_conversion_job(job: Job) -> None: if merged_required: audio_path = _build_output_path(audio_dir, job.original_filename, job.output_format) meta_for_sink = job.metadata_tags if job.metadata_tags else None - audio_sink = _open_audio_sink(audio_path, job, sink_stack, metadata=meta_for_sink) + audio_sink = sink_stack.enter_context( + open_audio_sink( + audio_path, + job.output_format, + metadata=meta_for_sink, + cancel_check=lambda: job.cancel_requested, + ) + ) subtitle_writer = _create_subtitle_writer(job, audio_path) job.result.audio_path = audio_path if subtitle_writer: @@ -638,11 +637,12 @@ def run_conversion_job(job: Job) -> None: f"{Path(job.original_filename).stem}_{_slugify(chapter_display_title, idx)}", job.separate_chapters_format, ) - chapter_sink = _open_audio_sink( - chapter_audio_path, - job, - chapter_sink_stack, - fmt=job.separate_chapters_format, + chapter_sink = chapter_sink_stack.enter_context( + open_audio_sink( + chapter_audio_path, + job.separate_chapters_format, + cancel_check=lambda: job.cancel_requested, + ) ) speak_heading = bool(heading_text) @@ -907,11 +907,12 @@ def run_conversion_job(job: Job) -> None: f"{Path(job.original_filename).stem}_outro", job.separate_chapters_format, ) - chapter_sink = _open_audio_sink( - outro_audio_path, - job, - outro_sink_stack, - fmt=job.separate_chapters_format, + chapter_sink = outro_sink_stack.enter_context( + open_audio_sink( + outro_audio_path, + job.separate_chapters_format, + cancel_check=lambda: job.cancel_requested, + ) ) outro_segments = emit_text( @@ -1148,50 +1149,6 @@ def _prepare_project_layout(job: Job, base_dir: Path) -> tuple[Path, Path, Path, ) -def _open_audio_sink( - path: Path, - job: Job, - stack: ExitStack, - *, - fmt: Optional[str] = None, - metadata: Optional[Dict[str, str]] = None, -) -> AudioSink: - ffmpeg_cache_root = get_internal_cache_path("ffmpeg") - platform_cache = os.path.join(ffmpeg_cache_root, sys.platform) - os.makedirs(platform_cache, exist_ok=True) - try: - import static_ffmpeg.run as static_ffmpeg_run # type: ignore - - static_ffmpeg_run.LOCK_FILE = os.path.join(ffmpeg_cache_root, "lock.file") - except Exception: - pass - - static_ffmpeg.add_paths(weak=True, download_dir=platform_cache) - fmt_value = (fmt or job.output_format).lower() - - if fmt_value in {"wav", "flac"}: - soundfile = stack.enter_context( - sf.SoundFile(path, mode="w", samplerate=SAMPLE_RATE, channels=1, format=fmt_value.upper()) - ) - return AudioSink(write=lambda data: soundfile.write(data)) - - cmd = _build_ffmpeg_command(path, fmt_value, metadata=metadata) - process = create_process(cmd, stdin=subprocess.PIPE, text=False) - - def _finalize() -> None: - if process.stdin and not process.stdin.closed: - process.stdin.close() - process.wait() - - stack.callback(_finalize) - - def _write(data: np.ndarray) -> None: - if job.cancel_requested or process.stdin is None: - return - process.stdin.write(data.tobytes()) # type: ignore[arg-type] - - return AudioSink(write=_write) - def _resolve_voice(pipeline, voice_spec: str, use_gpu: bool): if "*" in voice_spec: diff --git a/tests/test_domain_audio_sink.py b/tests/test_domain_audio_sink.py new file mode 100644 index 0000000..52ff28d --- /dev/null +++ b/tests/test_domain_audio_sink.py @@ -0,0 +1,189 @@ +"""Tests for domain/audio_sink.py""" + +import numpy as np +import soundfile as sf +from pathlib import Path +from unittest.mock import MagicMock, patch, call +import subprocess + +from abogen.domain.audio_sink import AudioSink, open_audio_sink, _ensure_ffmpeg + + +class TestAudioSinkDataclass: + def test_audio_sink_is_frozen(self): + sink = AudioSink(write=lambda d: None, close=lambda: None) + try: + sink.write = lambda d: None # type: ignore + except AttributeError: + pass # Expected — frozen=True + assert hasattr(sink, "write") + assert hasattr(sink, "close") + + def test_audio_sink_stores_callables(self): + write_fn = MagicMock() + close_fn = MagicMock() + sink = AudioSink(write=write_fn, close=close_fn) + assert sink.write is write_fn + assert sink.close is close_fn + + def test_audio_sink_write_callable(self): + calls = [] + sink = AudioSink(write=lambda d: calls.append(d), close=lambda: None) + data = np.zeros(100, dtype="float32") + sink.write(data) + assert len(calls) == 1 + np.testing.assert_array_equal(calls[0], data) + + def test_audio_sink_close_callable(self): + closed = [] + sink = AudioSink(write=lambda d: None, close=lambda: closed.append(True)) + sink.close() + assert closed == [True] + + +class TestOpenAudioSinkWav: + def test_wav_creates_file(self, tmp_path: Path): + out = tmp_path / "test.wav" + with open_audio_sink(out, "wav") as sink: + sink.write(np.zeros(100, dtype="float32")) + assert out.exists() + + def test_flac_creates_file(self, tmp_path: Path): + out = tmp_path / "test.flac" + with open_audio_sink(out, "flac") as sink: + sink.write(np.zeros(100, dtype="float32")) + assert out.exists() + + def test_wav_sink_writes_audio(self, tmp_path: Path): + out = tmp_path / "out.wav" + audio = np.random.uniform(-0.5, 0.5, 24000).astype("float32") # 1 second + with open_audio_sink(out, "wav") as sink: + sink.write(audio) + data, sr = sf.read(str(out)) + assert sr == 24000 + assert len(data) == 24000 + np.testing.assert_allclose(data, audio, atol=1e-4) + + def test_wav_sink_close_flushes(self, tmp_path: Path): + out = tmp_path / "flush.wav" + with open_audio_sink(out, "wav") as sink: + sink.write(np.ones(1000, dtype="float32")) + assert out.exists() + info = sf.info(str(out)) + assert info.samplerate == 24000 + assert info.channels == 1 + + def test_wav_sink_context_manager(self, tmp_path: Path): + out = tmp_path / "ctx.wav" + with open_audio_sink(out, "wav") as sink: + assert isinstance(sink, AudioSink) + sink.write(np.zeros(50, dtype="float32")) + assert out.exists() + + def test_wav_sink_multiple_writes(self, tmp_path: Path): + out = tmp_path / "multi.wav" + with open_audio_sink(out, "wav") as sink: + sink.write(np.ones(1000, dtype="float32")) + sink.write(np.ones(500, dtype="float32")) + data, _ = sf.read(str(out)) + assert len(data) == 1500 + + +class TestCancelCheck: + def test_cancel_check_skips_wav_writes(self, tmp_path: Path): + out = tmp_path / "cancelled.wav" + with open_audio_sink(out, "wav", cancel_check=lambda: True) as sink: + sink.write(np.ones(1000, dtype="float32")) + sink.write(np.ones(500, dtype="float32")) + data, _ = sf.read(str(out)) + assert len(data) == 0 + + def test_cancel_check_none_allows_writes(self, tmp_path: Path): + out = tmp_path / "ok.wav" + with open_audio_sink(out, "wav", cancel_check=None) as sink: + sink.write(np.ones(1000, dtype="float32")) + data, _ = sf.read(str(out)) + assert len(data) == 1000 + + +class TestUnsupportedFormat: + def test_unsupported_format_raises(self, tmp_path: Path): + out = tmp_path / "bad.xyz" + try: + with open_audio_sink(out, "xyz") as sink: + pass + assert False, "Should have raised" + except Exception: + pass # Expected + + +class TestOpenAudioSinkCompressed: + @patch("abogen.domain.audio_sink._ensure_ffmpeg") + @patch("abogen.domain.audio_sink.build_ffmpeg_command") + @patch("abogen.domain.audio_sink.subprocess.Popen") + def test_mp3_sink_returns_sink(self, mock_popen, mock_build, mock_ensure, tmp_path: Path): + mock_build.return_value = ["ffmpeg", "-y", "-i", "pipe:0", "out.mp3"] + proc = MagicMock() + proc.stdin = MagicMock() + proc.stdin.closed = False + proc.wait.return_value = 0 + mock_popen.return_value = proc + + out = tmp_path / "test.mp3" + sink = open_audio_sink(out, "mp3") + assert isinstance(sink, AudioSink) + sink.write(np.zeros(100, dtype="float32")) + assert proc.stdin.write.called + sink.close() + + @patch("abogen.domain.audio_sink._ensure_ffmpeg") + @patch("abogen.domain.audio_sink.build_ffmpeg_command") + @patch("abogen.domain.audio_sink.subprocess.Popen") + def test_cancel_check_skips_compressed_writes(self, mock_popen, mock_build, mock_ensure, tmp_path: Path): + mock_build.return_value = ["ffmpeg", "-y", "-i", "pipe:0", "out.mp3"] + proc = MagicMock() + proc.stdin = MagicMock() + proc.stdin.closed = False + proc.wait.return_value = 0 + mock_popen.return_value = proc + + sink = open_audio_sink(tmp_path / "c.mp3", "mp3", cancel_check=lambda: True) + sink.write(np.zeros(100, dtype="float32")) + assert not proc.stdin.write.called + sink.close() + + @patch("abogen.domain.audio_sink._ensure_ffmpeg") + @patch("abogen.domain.audio_sink.build_ffmpeg_command") + @patch("abogen.domain.audio_sink.subprocess.Popen") + def test_extra_ffmpeg_args_passed(self, mock_popen, mock_build, mock_ensure, tmp_path: Path): + mock_build.return_value = ["ffmpeg", "-y", "-header", "-i", "pipe:0", "out.mp3"] + proc = MagicMock() + proc.stdin = MagicMock() + proc.stdin.closed = False + proc.wait.return_value = 0 + mock_popen.return_value = proc + + sink = open_audio_sink( + tmp_path / "extra.mp3", + "mp3", + extra_ffmpeg_args=["-thread_queue_size", "32768"], + ) + args = mock_popen.call_args[0][0] + assert "-thread_queue_size" in args + assert "32768" in args + sink.close() + + @patch("abogen.domain.audio_sink._ensure_ffmpeg") + @patch("abogen.domain.audio_sink.build_ffmpeg_command") + @patch("abogen.domain.audio_sink.subprocess.Popen") + def test_metadata_passed_to_ffmpeg(self, mock_popen, mock_build, mock_ensure, tmp_path: Path): + mock_build.return_value = ["ffmpeg", "-y", "-i", "pipe:0", "out.opus"] + proc = MagicMock() + proc.stdin = MagicMock() + proc.stdin.closed = False + proc.wait.return_value = 0 + mock_popen.return_value = proc + + meta = {"title": "Test", "artist": "Author"} + open_audio_sink(tmp_path / "meta.opus", "opus", metadata=meta) + mock_build.assert_called_once_with(tmp_path / "meta.opus", "opus", metadata=meta)