lofi-app/tests/unit/test_stream_extractor.py

from unittest.mock import MagicMock, patch

import yt_dlp

from src.modules.stream_extractor import extract_audio_stream


class TestExtractAudioStream:
    def test_returns_hls_stream(self) -> None:
        mock_info = {
            "title": "Live Lofi Stream",
            "channel": "Lofi Girl",
            "duration": None,
            "live_status": "live",
            "formats": [
                {
                    "protocol": "m3u8",
                    "url": "https://manifest.hls.tv/pl.m3u8",
                    "format_note": "audio",
                    "acodec": "mp4a.40.2",
                }
            ],
        }

        mock_ydl = MagicMock()
        mock_ydl.extract_info.return_value = mock_info
        mock_ydl.__enter__ = MagicMock(return_value=mock_ydl)
        mock_ydl.__exit__ = MagicMock(return_value=False)

        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")

        assert result is not None
        assert result["videoId"] == "dQw4w9WgXcQ"
        assert result["url"] == "https://manifest.hls.tv/pl.m3u8"
        assert result["streamType"] == "hls"
        assert result["isLive"] is True

    def test_returns_direct_stream(self) -> None:
        mock_info = {
            "title": "Lofi Beats",
            "channel": "Chillhop",
            "duration": 3600,
            "live_status": "none",
            "formats": [
                {
                    "protocol": "https",
                    "url": "https://rr3---sn-5hu.googlevideo.com/audio.mp4",
                    "format_note": "audio only",
                    "acodec": "mp4a.40.2",
                }
            ],
        }

        mock_ydl = MagicMock()
        mock_ydl.extract_info.return_value = mock_info
        mock_ydl.__enter__ = MagicMock(return_value=mock_ydl)
        mock_ydl.__exit__ = MagicMock(return_value=False)

        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")

        assert result is not None
        assert result["streamType"] == "direct"
        assert result["isLive"] is False

    def _mock_ydl(self, extract_info_return=None, side_effect=None):
        mock_ydl = MagicMock()
        if side_effect is not None:
            mock_ydl.extract_info.side_effect = side_effect
        else:
            mock_ydl.extract_info.return_value = extract_info_return
        mock_ydl.__enter__ = MagicMock(return_value=mock_ydl)
        mock_ydl.__exit__ = MagicMock(return_value=False)
        return mock_ydl

    def test_returns_none_when_no_info(self) -> None:
        mock_ydl = self._mock_ydl(extract_info_return=None)
        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")
        assert result is None

    def test_returns_none_when_no_formats(self) -> None:
        mock_ydl = self._mock_ydl(extract_info_return={"formats": []})
        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")
        assert result is None

    def test_returns_none_when_video_only_formats(self) -> None:
        mock_ydl = self._mock_ydl(
            extract_info_return={
                "formats": [
                    {
                        "protocol": "https",
                        "url": "https://video.mp4",
                        "format_note": "video only",
                        "acodec": "none",
                        "vcodec": "avc1",
                    }
                ]
            }
        )
        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")
        assert result is None

    def test_prefers_hls_over_direct(self) -> None:
        mock_ydl = self._mock_ydl(
            extract_info_return={
                "title": "Stream",
                "channel": "Channel",
                "duration": None,
                "live_status": "live",
                "formats": [
                    {
                        "protocol": "https",
                        "url": "https://direct-audio.mp4",
                        "format_note": "audio",
                        "acodec": "mp4a.40.2",
                    },
                    {
                        "protocol": "m3u8",
                        "url": "https://hls-audio.m3u8",
                        "format_note": "audio",
                        "acodec": "mp4a.40.2",
                    },
                ],
            }
        )
        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")
        assert result["streamType"] == "hls"
        assert result["url"] == "https://hls-audio.m3u8"

    def test_handles_download_error(self) -> None:
        mock_ydl = self._mock_ydl(
            side_effect=yt_dlp.utils.DownloadError("Video unavailable")
        )
        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")
        assert result is None

    def test_handles_unexpected_error(self) -> None:
        mock_ydl = self._mock_ydl(side_effect=Exception("Unexpected error"))
        with patch(
            "src.modules.stream_extractor.yt_dlp.YoutubeDL", return_value=mock_ydl
        ):
            result = extract_audio_stream("dQw4w9WgXcQ")
        assert result is None