test_write_annotations.py.disabled 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. #!/usr/bin/env python3
  2. # Allow direct execution
  3. import os
  4. import sys
  5. import unittest
  6. sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  7. import xml.etree.ElementTree
  8. import yt_dlp.extractor
  9. import yt_dlp.YoutubeDL
  10. from test.helper import get_params, is_download_test, try_rm
  11. class YoutubeDL(yt_dlp.YoutubeDL):
  12. def __init__(self, *args, **kwargs):
  13. super().__init__(*args, **kwargs)
  14. self.to_stderr = self.to_screen
  15. params = get_params({
  16. 'writeannotations': True,
  17. 'skip_download': True,
  18. 'writeinfojson': False,
  19. 'format': 'flv',
  20. })
  21. TEST_ID = 'gr51aVj-mLg'
  22. ANNOTATIONS_FILE = TEST_ID + '.annotations.xml'
  23. EXPECTED_ANNOTATIONS = ['Speech bubble', 'Note', 'Title', 'Spotlight', 'Label']
  24. @is_download_test
  25. class TestAnnotations(unittest.TestCase):
  26. def setUp(self):
  27. # Clear old files
  28. self.tearDown()
  29. def test_info_json(self):
  30. expected = list(EXPECTED_ANNOTATIONS) # Two annotations could have the same text.
  31. ie = yt_dlp.extractor.YoutubeIE()
  32. ydl = YoutubeDL(params)
  33. ydl.add_info_extractor(ie)
  34. ydl.download([TEST_ID])
  35. self.assertTrue(os.path.exists(ANNOTATIONS_FILE))
  36. annoxml = None
  37. with open(ANNOTATIONS_FILE, encoding='utf-8') as annof:
  38. annoxml = xml.etree.ElementTree.parse(annof)
  39. self.assertTrue(annoxml is not None, 'Failed to parse annotations XML')
  40. root = annoxml.getroot()
  41. self.assertEqual(root.tag, 'document')
  42. annotationsTag = root.find('annotations')
  43. self.assertEqual(annotationsTag.tag, 'annotations')
  44. annotations = annotationsTag.findall('annotation')
  45. # Not all the annotations have TEXT children and the annotations are returned unsorted.
  46. for a in annotations:
  47. self.assertEqual(a.tag, 'annotation')
  48. if a.get('type') == 'text':
  49. textTag = a.find('TEXT')
  50. text = textTag.text
  51. self.assertTrue(text in expected) # assertIn only added in python 2.7
  52. # remove the first occurrence, there could be more than one annotation with the same text
  53. expected.remove(text)
  54. # We should have seen (and removed) all the expected annotation texts.
  55. self.assertEqual(len(expected), 0, 'Not all expected annotations were found.')
  56. def tearDown(self):
  57. try_rm(ANNOTATIONS_FILE)
  58. if __name__ == '__main__':
  59. unittest.main()