Redid errors for cookies, improved testing coverage

This commit is contained in:
danielcliu 2020-01-30 21:53:18 -08:00
parent 49ccba7f95
commit 31b8f4a179
4 changed files with 38 additions and 19 deletions

View File

@ -8,4 +8,6 @@ from ._errors import (
NotTranslatable,
TranslationLanguageNotAvailable,
NoTranscriptAvailable,
CookiePathInvalid,
CookiesInvalid
)

View File

@ -1,8 +1,10 @@
import requests
try:
import http.cookiejar as cookiejar
CookieLoadError = (FileNotFoundError, cookiejar.LoadError)
except ImportError:
import cookielib as cookiejar
CookieLoadError = IOError
from ._transcripts import TranscriptListFetcher
@ -63,7 +65,7 @@ class YouTubeTranscriptApi():
"""
with requests.Session() as http_client:
if cookies:
http_client.cookies = cls.load_cookies(cookies)
http_client.cookies = cls._load_cookies(cookies, video_id)
http_client.proxies = proxies if proxies else {}
return TranscriptListFetcher(http_client).fetch(video_id)
@ -126,15 +128,13 @@ class YouTubeTranscriptApi():
return cls.list_transcripts(video_id, proxies, cookies).find_transcript(languages).fetch()
@classmethod
def load_cookies(cls, cookies):
cj = {}
def _load_cookies(cls, cookies, video_id):
cookie_jar = {}
try:
cj = cookiejar.MozillaCookieJar()
cj.load(cookies)
except IOError as e:
raise CookiePathInvalid
except FileNotFoundError as e:
raise CookiePathInvalid
if not cj:
raise CookiesInvalid
return cj
cookie_jar = cookiejar.MozillaCookieJar()
cookie_jar.load(cookies)
except CookieLoadError:
raise CookiePathInvalid(video_id)
if not cookie_jar:
raise CookiesInvalid(video_id)
return cookie_jar

View File

@ -56,7 +56,7 @@ class TranslationLanguageNotAvailable(CouldNotRetrieveTranscript):
class CookiePathInvalid(CouldNotRetrieveTranscript):
CAUSE_MESSAGE = 'Path to cookie file was not valid'
CAUSE_MESSAGE = 'The provided cookie file was unable to be loaded'
class CookiesInvalid(CouldNotRetrieveTranscript):

View File

@ -15,6 +15,8 @@ from youtube_transcript_api import (
NoTranscriptAvailable,
NotTranslatable,
TranslationLanguageNotAvailable,
CookiePathInvalid,
CookiesInvalid
)
@ -166,6 +168,20 @@ class TestYouTubeTranscriptApi(TestCase):
{'text': 'just something shorter, I made up for testing', 'start': 5.7, 'duration': 3.239}
]
)
def test_get_transcript__with_cookies(self):
dirname, filename = os.path.split(os.path.abspath(__file__))
cookies = dirname + '/example_cookies.txt'
transcript = YouTubeTranscriptApi.get_transcript('GJLlxj_dtq8', cookies=cookies)
self.assertEqual(
transcript,
[
{'text': 'Hey, this is just a test', 'start': 0.0, 'duration': 1.54},
{'text': 'this is not the original transcript', 'start': 1.54, 'duration': 4.16},
{'text': 'just something shorter, I made up for testing', 'start': 5.7, 'duration': 3.239}
]
)
@patch('youtube_transcript_api.YouTubeTranscriptApi.get_transcript')
def test_get_transcripts(self, mock_get_transcript):
@ -209,15 +225,16 @@ class TestYouTubeTranscriptApi(TestCase):
def test_load_cookies(self):
dirname, filename = os.path.split(os.path.abspath(__file__))
cookies = dirname + '/example_cookies.txt'
session_cookies = YouTubeTranscriptApi.load_cookies(cookies)
session_cookies = YouTubeTranscriptApi._load_cookies(cookies, 'GJLlxj_dtq8')
self.assertEqual({'TEST_FIELD': 'TEST_VALUE'}, requests.utils.dict_from_cookiejar(session_cookies))
def test_load_cookies__bad_files(self):
def test_load_cookies__bad_file_path(self):
bad_cookies = 'nonexistent_cookies.txt'
with self.assertRaises(Exception):
YouTubeTranscriptApi.load_cookies(bad_cookies)
with self.assertRaises(CookiePathInvalid):
YouTubeTranscriptApi._load_cookies(bad_cookies, 'GJLlxj_dtq8')
def test_load_cookies__no_valid_cookies(self):
dirname, filename = os.path.split(os.path.abspath(__file__))
expired_cookies = dirname + '/expired_example_cookies.txt'
with self.assertRaises(Exception):
YouTubeTranscriptApi.load_cookies(expired_cookies)
with self.assertRaises(CookiesInvalid):
YouTubeTranscriptApi._load_cookies(expired_cookies, 'GJLlxj_dtq8')