Add support for cookie file to media_extractor (#104973)

* media_extractor comp -> add support for cookie file

* Update __init__.py

* Update __init__.py

* fixed pr request, added test

* update test with valid cookie file

* move cookies to subdirectory

* use pathlib instead of os.path
This commit is contained in:
DiamondDrake 2023-12-30 08:34:50 -05:00 committed by GitHub
parent faa2129e96
commit 8140036d2e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 74 additions and 1 deletions

View File

@ -1,6 +1,7 @@
"""Decorator service for the media_player.play_media service.""" """Decorator service for the media_player.play_media service."""
from collections.abc import Callable from collections.abc import Callable
import logging import logging
from pathlib import Path
from typing import Any, cast from typing import Any, cast
import voluptuous as vol import voluptuous as vol
@ -106,7 +107,20 @@ class MediaExtractor:
def get_stream_selector(self) -> Callable[[str], str]: def get_stream_selector(self) -> Callable[[str], str]:
"""Return format selector for the media URL.""" """Return format selector for the media URL."""
ydl = YoutubeDL({"quiet": True, "logger": _LOGGER}) cookies_file = Path(
self.hass.config.config_dir, "media_extractor", "cookies.txt"
)
ydl_params = {"quiet": True, "logger": _LOGGER}
if cookies_file.exists():
ydl_params["cookiefile"] = str(cookies_file)
_LOGGER.debug(
"Media extractor loaded cookies file from: %s", str(cookies_file)
)
else:
_LOGGER.debug(
"Media extractor didn't find cookies file at: %s", str(cookies_file)
)
ydl = YoutubeDL(ydl_params)
try: try:
all_media = ydl.extract_info(self.get_media_url(), process=False) all_media = ydl.extract_info(self.get_media_url(), process=False)

View File

@ -1,4 +1,6 @@
"""The tests for Media Extractor integration.""" """The tests for Media Extractor integration."""
import os
import os.path
from typing import Any from typing import Any
from unittest.mock import patch from unittest.mock import patch
@ -209,3 +211,60 @@ async def test_query_error(
await hass.async_block_till_done() await hass.async_block_till_done()
assert len(calls) == 0 assert len(calls) == 0
async def test_cookiefile_detection(
hass: HomeAssistant,
mock_youtube_dl: MockYoutubeDL,
empty_media_extractor_config: dict[str, Any],
calls: list[ServiceCall],
snapshot: SnapshotAssertion,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test cookie file detection."""
await async_setup_component(hass, DOMAIN, empty_media_extractor_config)
await hass.async_block_till_done()
cookies_dir = os.path.join(hass.config.config_dir, "media_extractor")
cookies_file = os.path.join(cookies_dir, "cookies.txt")
if not os.path.exists(cookies_dir):
os.makedirs(cookies_dir)
f = open(cookies_file, "w+", encoding="utf-8")
f.write(
"""# Netscape HTTP Cookie File
.youtube.com TRUE / TRUE 1701708706 GPS 1
"""
)
f.close()
await hass.services.async_call(
DOMAIN,
SERVICE_PLAY_MEDIA,
{
"entity_id": "media_player.bedroom",
"media_content_type": "VIDEO",
"media_content_id": YOUTUBE_PLAYLIST,
},
)
await hass.async_block_till_done()
assert "Media extractor loaded cookies file" in caplog.text
os.remove(cookies_file)
await hass.services.async_call(
DOMAIN,
SERVICE_PLAY_MEDIA,
{
"entity_id": "media_player.bedroom",
"media_content_type": "VIDEO",
"media_content_id": YOUTUBE_PLAYLIST,
},
)
await hass.async_block_till_done()
assert "Media extractor didn't find cookies file" in caplog.text