From 869cf6b4a92fdadc15d06c78b37b257effb797a0 Mon Sep 17 00:00:00 2001
From: adavyas <adavyasharma@gmail.com>
Date: Wed, 8 Apr 2026 19:10:24 -0400
Subject: [PATCH] Fix audio chunking edge cases

---
 src/utils/files.py                   | 16 ++++++-
 tests/utils/test_audio_processing.py | 65 ++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/src/utils/files.py b/src/utils/files.py
index 4f1f53cb..c9c5de3a 100644
--- a/src/utils/files.py
+++ b/src/utils/files.py
@@ -219,13 +219,14 @@ class AudioProcessor:
             raise ValidationException("Audio upload is empty")
 
         suffix = self.get_output_suffix(filename, content_type)
+        normalized_filename = self.ensure_audio_filename(filename, suffix)
         duration_seconds = self._probe_audio_duration_seconds(content, suffix)
 
         if (
             len(content) <= settings.AUDIO.MAX_CHUNK_BYTES
             and duration_seconds <= settings.AUDIO.MAX_CHUNK_DURATION_SECONDS
         ):
-            return [AudioSegment(index=0, filename=filename, content=content)]
+            return [AudioSegment(index=0, filename=normalized_filename, content=content)]
 
         segment_count = self._estimate_initial_segment_count(
             duration_seconds=duration_seconds,
@@ -253,6 +254,13 @@ class AudioProcessor:
                     return segments
 
                 if duration_seconds / segment_count <= 1.0:
+                    if suffix == ".wav":
+                        suffix = ".mp3"
+                        segment_count = self._estimate_initial_segment_count(
+                            duration_seconds=duration_seconds,
+                            suffix=suffix,
+                        )
+                        continue
                     raise FileProcessingError(
                         "Audio segmentation could not satisfy max chunk size"
                     )
@@ -347,6 +355,12 @@ class AudioProcessor:
             return ".wav"
         return ".mp3"
 
+    def ensure_audio_filename(self, filename: str, suffix: str) -> str:
+        path = Path(filename)
+        if path.suffix.lower() == suffix:
+            return filename
+        return f"{filename}{suffix}"
+
     def _probe_audio_duration_seconds(self, content: bytes, suffix: str) -> float:
         with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as temp_file:
             temp_file.write(content)
diff --git a/tests/utils/test_audio_processing.py b/tests/utils/test_audio_processing.py
index 2aae23a1..d1965e79 100644
--- a/tests/utils/test_audio_processing.py
+++ b/tests/utils/test_audio_processing.py
@@ -342,6 +342,71 @@ def test_split_audio_segments_splits_long_wav_when_duration_exceeds_limit():
     assert all(segment.content for segment in segments)
 
 
+def test_split_audio_segments_normalizes_small_mime_only_audio_filename():
+    processor = AudioProcessor()
+
+    original_duration_limit = settings.AUDIO.MAX_CHUNK_DURATION_SECONDS
+    original_chunk_bytes = settings.AUDIO.MAX_CHUNK_BYTES
+    settings.AUDIO.MAX_CHUNK_DURATION_SECONDS = 60
+    settings.AUDIO.MAX_CHUNK_BYTES = 1024
+    try:
+        with patch.object(
+            processor,
+            "_probe_audio_duration_seconds",
+            return_value=0.5,
+        ):
+            segments = processor.split_audio_segments(
+                b"audio-bytes",
+                filename="blob",
+                content_type="audio/mpeg",
+            )
+    finally:
+        settings.AUDIO.MAX_CHUNK_DURATION_SECONDS = original_duration_limit
+        settings.AUDIO.MAX_CHUNK_BYTES = original_chunk_bytes
+
+    assert len(segments) == 1
+    assert segments[0].filename == "blob.mp3"
+
+
+def test_split_audio_segments_falls_back_to_mp3_when_wav_floor_exceeds_byte_limit():
+    processor = AudioProcessor()
+
+    original_duration_limit = settings.AUDIO.MAX_CHUNK_DURATION_SECONDS
+    original_chunk_bytes = settings.AUDIO.MAX_CHUNK_BYTES
+    settings.AUDIO.MAX_CHUNK_DURATION_SECONDS = 60
+    settings.AUDIO.MAX_CHUNK_BYTES = 200_000
+    build_segments = [
+        [AudioSegment(index=0, filename="segment_000.wav", content=b"x" * 200_001)],
+        [AudioSegment(index=0, filename="segment_000.mp3", content=b"x" * 10)],
+    ]
+    try:
+        with (
+            patch.object(
+                processor,
+                "_probe_audio_duration_seconds",
+                return_value=1.0,
+            ),
+            patch.object(
+                processor,
+                "_build_audio_segments",
+                side_effect=build_segments,
+            ) as mock_build,
+        ):
+            segments = processor.split_audio_segments(
+                b"x" * 200_001,
+                filename="clip.wav",
+                content_type="audio/wav",
+            )
+    finally:
+        settings.AUDIO.MAX_CHUNK_DURATION_SECONDS = original_duration_limit
+        settings.AUDIO.MAX_CHUNK_BYTES = original_chunk_bytes
+
+    assert len(segments) == 1
+    assert segments[0].filename.endswith(".mp3")
+    assert mock_build.call_args_list[0].kwargs["suffix"] == ".wav"
+    assert mock_build.call_args_list[1].kwargs["suffix"] == ".mp3"
+
+
 def test_split_audio_segments_raises_validation_for_invalid_audio_bytes():
     processor = AudioProcessor()