X-Git-Url: https://git.rapsys.eu/youtubedl/blobdiff_plain/b238854ce845f3796daac74edab2e8a373e8ba1a..92816565411fb77b1c3d9f49e414f380bc131a00:/test/test_write_annotations.py

diff --git a/test/test_write_annotations.py b/test/test_write_annotations.py
index eac53b2..41abdfe 100644
--- a/test/test_write_annotations.py
+++ b/test/test_write_annotations.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 # coding: utf-8
+from __future__ import unicode_literals
 
 # Allow direct execution
 import os
@@ -23,6 +24,7 @@ class YoutubeDL(youtube_dl.YoutubeDL):
         super(YoutubeDL, self).__init__(*args, **kwargs)
         self.to_stderr = self.to_screen
 
+
 params = get_params({
     'writeannotations': True,
     'skip_download': True,
@@ -31,19 +33,18 @@ params = get_params({
 })
 
 
-
 TEST_ID = 'gr51aVj-mLg'
-ANNOTATIONS_FILE = TEST_ID + '.flv.annotations.xml'
+ANNOTATIONS_FILE = TEST_ID + '.annotations.xml'
 EXPECTED_ANNOTATIONS = ['Speech bubble', 'Note', 'Title', 'Spotlight', 'Label']
 
+
 class TestAnnotations(unittest.TestCase):
     def setUp(self):
         # Clear old files
         self.tearDown()
 
-
     def test_info_json(self):
-        expected = list(EXPECTED_ANNOTATIONS) #Two annotations could have the same text.
+        expected = list(EXPECTED_ANNOTATIONS)  # Two annotations could have the same text.
         ie = youtube_dl.extractor.YoutubeIE()
         ydl = YoutubeDL(params)
         ydl.add_info_extractor(ie)
@@ -51,7 +52,7 @@ class TestAnnotations(unittest.TestCase):
         self.assertTrue(os.path.exists(ANNOTATIONS_FILE))
         annoxml = None
         with io.open(ANNOTATIONS_FILE, 'r', encoding='utf-8') as annof:
-                annoxml = xml.etree.ElementTree.parse(annof)
+            annoxml = xml.etree.ElementTree.parse(annof)
         self.assertTrue(annoxml is not None, 'Failed to parse annotations XML')
         root = annoxml.getroot()
         self.assertEqual(root.tag, 'document')
@@ -59,21 +60,21 @@ class TestAnnotations(unittest.TestCase):
         self.assertEqual(annotationsTag.tag, 'annotations')
         annotations = annotationsTag.findall('annotation')
 
-        #Not all the annotations have TEXT children and the annotations are returned unsorted.
+        # Not all the annotations have TEXT children and the annotations are returned unsorted.
         for a in annotations:
-                self.assertEqual(a.tag, 'annotation')
-                if a.get('type') == 'text':
-                        textTag = a.find('TEXT')
-                        text = textTag.text
-                        self.assertTrue(text in expected) #assertIn only added in python 2.7
-                        #remove the first occurance, there could be more than one annotation with the same text
-                        expected.remove(text)
-        #We should have seen (and removed) all the expected annotation texts.
+            self.assertEqual(a.tag, 'annotation')
+            if a.get('type') == 'text':
+                textTag = a.find('TEXT')
+                text = textTag.text
+                self.assertTrue(text in expected)  # assertIn only added in python 2.7
+                # remove the first occurrence, there could be more than one annotation with the same text
+                expected.remove(text)
+        # We should have seen (and removed) all the expected annotation texts.
         self.assertEqual(len(expected), 0, 'Not all expected annotations were found.')
-        
 
     def tearDown(self):
         try_rm(ANNOTATIONS_FILE)
 
+
 if __name__ == '__main__':
     unittest.main()