[test] set up initial sabr processor tests

test/test_sabr/test_processor.py

@@ -0,0 +1,252 @@
+import pytest
+from unittest.mock import MagicMock
+
+from yt_dlp.extractor.youtube._streaming.sabr.processor import SabrProcessor
+from yt_dlp.extractor.youtube._streaming.sabr.models import (
+    AudioSelector,
+    VideoSelector,
+    CaptionSelector,
+)
+from yt_dlp.extractor.youtube._proto.videostreaming import FormatId
+from yt_dlp.extractor.youtube._proto.innertube import ClientInfo
+
+
+@pytest.fixture
+def logger():
+    return MagicMock()
+
+
+@pytest.fixture
+def client_info():
+    return ClientInfo()
+
+
+@pytest.fixture
+def base_args(logger, client_info):
+    return {
+        'logger': logger,
+        'client_info': client_info,
+        'video_playback_ustreamer_config': 'dGVzdA==',
+    }
+
+
+def make_selector(selector_type, *, discard_media=False, format_ids=None):
+    if selector_type == 'audio':
+        return AudioSelector(
+            display_name='audio',
+            format_ids=format_ids or [FormatId(itag=140)],
+            discard_media=discard_media,
+        )
+    elif selector_type == 'video':
+        return VideoSelector(
+            display_name='video',
+            format_ids=format_ids or [FormatId(itag=248)],
+            discard_media=discard_media,
+        )
+    elif selector_type == 'caption':
+        return CaptionSelector(
+            display_name='caption',
+            format_ids=format_ids or [FormatId(itag=386)],
+            discard_media=discard_media,
+        )
+    raise ValueError(f'Unknown selector_type: {selector_type}')
+
+
+def selector_factory(selector_type, *, discard_media=False, format_ids=None):
+    def factory():
+        return make_selector(selector_type, discard_media=discard_media, format_ids=format_ids)
+    return factory
+
+
+class TestSabrProcessorInitialization:
+    @pytest.mark.parametrize(
+        'audio_sel,video_sel,caption_sel,expected_bitfield',
+        [
+            # audio+video
+            (selector_factory('audio'), selector_factory('video'), None, 0),
+            # audio+video+caption(discard)
+            (
+                selector_factory('audio'),
+                selector_factory('video'),
+                selector_factory('caption', discard_media=True),
+                0,
+            ),
+            # audio only
+            (selector_factory('audio'), None, None, 1),
+            # audio only (video+caption manual discard)
+            (
+                selector_factory('audio'),
+                selector_factory('video', discard_media=True),
+                selector_factory('caption', discard_media=True),
+                1,
+            ),
+            # audio+video+caption
+            (
+                selector_factory('audio'),
+                selector_factory('video'),
+                selector_factory('caption'),
+                7,
+            ),
+            # video only
+            (None, selector_factory('video'), None, 0),
+            # video only (audio+caption manual discard)
+            (
+                selector_factory('audio', discard_media=True),
+                selector_factory('video'),
+                selector_factory('caption', discard_media=True),
+                0,
+            ),
+            # caption only
+            (None, None, selector_factory('caption'), 7),
+            # caption only (audio+video manual discard)
+            (
+                selector_factory('audio', discard_media=True),
+                selector_factory('video', discard_media=True),
+                selector_factory('caption'),
+                7,
+            ),
+        ],
+    )
+    def test_client_abr_state_bitfield(
+        self, base_args, audio_sel, video_sel, caption_sel, expected_bitfield,
+    ):
+        processor = SabrProcessor(
+            **base_args,
+            audio_selection=audio_sel() if audio_sel else None,
+            video_selection=video_sel() if video_sel else None,
+            caption_selection=caption_sel() if caption_sel else None,
+        )
+        assert processor.client_abr_state.enabled_track_types_bitfield == expected_bitfield
+
+    @pytest.mark.parametrize(
+        'audio_sel,video_sel,caption_sel,expected_audio_ids,expected_video_ids,expected_caption_ids',
+        [
+            # audio+video
+            (
+                selector_factory('audio'), selector_factory('video'), None,
+                [FormatId(itag=140)], [FormatId(itag=248)], [],
+            ),
+            # audio only
+            (
+                selector_factory('audio'), None, None,
+                [FormatId(itag=140)], [], [],
+            ),
+            # video only
+            (
+                None, selector_factory('video'), None,
+                [], [FormatId(itag=248)], [],
+            ),
+            # caption only
+            (
+                None, None, selector_factory('caption'),
+                [], [], [FormatId(itag=386)],
+            ),
+            # audio+video+caption
+            (
+                selector_factory('audio'), selector_factory('video'), selector_factory('caption'),
+                [FormatId(itag=140)], [FormatId(itag=248)], [FormatId(itag=386)],
+            ),
+            # multiple ids
+            (
+                selector_factory('audio', format_ids=[FormatId(itag=140), FormatId(itag=141)]),
+                selector_factory('video', format_ids=[FormatId(itag=248), FormatId(itag=249)]),
+                selector_factory('caption', format_ids=[FormatId(itag=386), FormatId(itag=387)]),
+                [FormatId(itag=140), FormatId(itag=141)],
+                [FormatId(itag=248), FormatId(itag=249)],
+                [FormatId(itag=386), FormatId(itag=387)],
+            ),
+        ],
+    )
+    def test_selected_format_ids(
+        self, base_args, audio_sel, video_sel, caption_sel,
+        expected_audio_ids, expected_video_ids, expected_caption_ids,
+    ):
+        processor = SabrProcessor(
+            **base_args,
+            audio_selection=audio_sel() if audio_sel else None,
+            video_selection=video_sel() if video_sel else None,
+            caption_selection=caption_sel() if caption_sel else None,
+        )
+        assert processor.selected_audio_format_ids == expected_audio_ids
+        assert processor.selected_video_format_ids == expected_video_ids
+        assert processor.selected_caption_format_ids == expected_caption_ids
+
+    @pytest.mark.parametrize(
+        'start_time_ms,expected',
+        [
+            (None, 0),
+            (0, 0),
+            (12345, 12345),
+        ],
+    )
+    def test_start_time_ms_initialization(self, base_args, start_time_ms, expected):
+        processor = SabrProcessor(
+            **base_args,
+            start_time_ms=start_time_ms,
+        )
+        assert processor.start_time_ms == expected
+        assert processor.client_abr_state.player_time_ms == expected
+
+    @pytest.mark.parametrize('invalid_start_time_ms', [-1, -100])
+    def test_start_time_ms_invalid(self, base_args, invalid_start_time_ms):
+        with pytest.raises(ValueError, match='start_time_ms must be greater than or equal to 0'):
+            SabrProcessor(
+                **base_args,
+                audio_selection=selector_factory('audio')(),
+                video_selection=selector_factory('video')(),
+                caption_selection=None,
+                start_time_ms=invalid_start_time_ms,
+            )
+
+    @pytest.mark.parametrize(
+        'duration_sec,tolerance_ms',
+        [
+            (10, 4999),
+            (10, 0),
+        ],
+    )
+    def test_live_segment_target_duration_tolerance_ms_valid(self, base_args, duration_sec, tolerance_ms):
+        # Should not raise
+        SabrProcessor(
+            **base_args,
+            live_segment_target_duration_sec=duration_sec,
+            live_segment_target_duration_tolerance_ms=tolerance_ms,
+        )
+
+    @pytest.mark.parametrize(
+        'duration_sec,tolerance_ms',
+        [
+            (10, 5000),  # exactly half
+            (10, 6000),  # more than half
+        ],
+    )
+    def test_live_segment_target_duration_tolerance_ms_validation(self, base_args, duration_sec, tolerance_ms):
+        with pytest.raises(ValueError, match='live_segment_target_duration_tolerance_ms must be less than'):
+            SabrProcessor(
+                **base_args,
+                live_segment_target_duration_sec=duration_sec,
+                live_segment_target_duration_tolerance_ms=tolerance_ms,
+            )
+
+    def test_defaults(self, base_args):
+        processor = SabrProcessor(**base_args)
+        assert processor.live_segment_target_duration_sec == 5
+        assert processor.live_segment_target_duration_tolerance_ms == 100
+        assert processor.start_time_ms == 0
+        assert processor.live_end_segment_tolerance == 10
+        assert processor.post_live is False
+
+    def test_override_defaults(self, base_args):
+        processor = SabrProcessor(
+            **base_args,
+            live_segment_target_duration_sec=8,
+            live_segment_target_duration_tolerance_ms=42,
+            start_time_ms=123,
+            live_end_segment_tolerance=3,
+            post_live=True,
+        )
+        assert processor.live_segment_target_duration_sec == 8
+        assert processor.live_segment_target_duration_tolerance_ms == 42
+        assert processor.start_time_ms == 123
+        assert processor.live_end_segment_tolerance == 3
+        assert processor.post_live is True

yt_dlp/extractor/youtube/_streaming/sabr/processor.py

@@ -119,6 +119,7 @@ def __init__(
         if self.start_time_ms < 0:
             raise ValueError('start_time_ms must be greater than or equal to 0')
 
+        # TODO: move to SabrStream
         self.live_end_wait_sec = live_end_wait_sec or max(10, 3 * self.live_segment_target_duration_sec)
         self.live_end_segment_tolerance = live_end_segment_tolerance or 10
         self.post_live = post_live
@@ -157,25 +158,28 @@ def is_live(self, value: bool):
         self._is_live = value
 
     def _initialize_cabr_state(self):
-        enabled_track_types_bitfield = 0  # Audio+Video
+        # SABR supports: audio+video, audio+video+captions or audio-only.
+        # For the other cases, we'll mark the tracks to be discarded (and fully buffered on initialization)
+
         if not self._video_format_selector:
-            enabled_track_types_bitfield = 1  # Audio only
             self._video_format_selector = VideoSelector(display_name='video_ignore', discard_media=True)
 
-        if self._caption_format_selector:
-            # SABR does not support caption-only or audio+captions only - can only get audio+video with captions
-            # If audio or video is not selected, the tracks will be initialized but marked as buffered.
-            enabled_track_types_bitfield = 7
-
-        # SABR does not support video-only, so we need to discard the audio track received.
-        # We need a selector as the server sometimes does not like it
-        # if we haven't initialized an audio format (e.g. livestreams).
         if not self._audio_format_selector:
             self._audio_format_selector = AudioSelector(display_name='audio_ignore', discard_media=True)
 
         if not self._caption_format_selector:
             self._caption_format_selector = CaptionSelector(display_name='caption_ignore', discard_media=True)
 
+        enabled_track_types_bitfield = 0  # Audio+Video
+
+        if self._video_format_selector.discard_media:
+            enabled_track_types_bitfield = 1  # Audio only
+
+        if not self._caption_format_selector.discard_media:
+            # SABR does not support caption-only or audio+captions only - can only get audio+video with captions
+            # If audio or video is not selected, the tracks will be initialized but marked as buffered.
+            enabled_track_types_bitfield = 7
+
         self.selected_audio_format_ids = self._audio_format_selector.format_ids
         self.selected_video_format_ids = self._video_format_selector.format_ids
         self.selected_caption_format_ids = self._caption_format_selector.format_ids