[cleanup] Misc cleanup

[yt-dlp.git] / test / helper.py

diff --git a/test/helper.py b/test/helper.py
index f45818b0f124d7e5ecaf70818176bc27fa46be5d..5c0e645f956b516d43439a6db59acdbcf052cd8f 100644 (file)
--- a/test/helper.py
+++ b/test/helper.py
@@ -10,23 +10,31 @@
 import ssl
 import sys
 
-import youtube_dlc.extractor
-from youtube_dlc import YoutubeDL
-from youtube_dlc.compat import (
+import yt_dlp.extractor
+from yt_dlp import YoutubeDL
+from yt_dlp.compat import (
     compat_os_name,
     compat_str,
 )
-from youtube_dlc.utils import (
+from yt_dlp.utils import (
     preferredencoding,
     write_string,
 )
 
 
+if 'pytest' in sys.modules:
+    import pytest
+    is_download_test = pytest.mark.download
+else:
+    def is_download_test(testClass):
+        return testClass
+
+
 def get_params(override=None):
     PARAMETERS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
-                                   "parameters.json")
+                                   'parameters.json')
     LOCAL_PARAMETERS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
-                                         "local_parameters.json")
+                                         'local_parameters.json')
     with io.open(PARAMETERS_FILE, encoding='utf-8') as pf:
         parameters = json.load(pf)
     if os.path.exists(LOCAL_PARAMETERS_FILE):
@@ -90,7 +98,7 @@ def report_warning(self, message):
 
 
 def gettestcases(include_onlymatching=False):
-    for ie in youtube_dlc.extractor.gen_extractors():
+    for ie in yt_dlp.extractor.gen_extractors():
         for tc in ie.get_testcases(include_onlymatching):
             yield tc
 
@@ -190,7 +198,10 @@ def expect_info_dict(self, got_dict, expected_dict):
     expect_dict(self, got_dict, expected_dict)
     # Check for the presence of mandatory fields
     if got_dict.get('_type') not in ('playlist', 'multi_video'):
-        for key in ('id', 'url', 'title', 'ext'):
+        mandatory_fields = ['id', 'title']
+        if expected_dict.get('ext'):
+            mandatory_fields.extend(('url', 'ext'))
+        for key in mandatory_fields:
             self.assertTrue(got_dict.get(key), 'Missing mandatory field %s' % key)
     # Check for mandatory fields that are automatically set by YoutubeDL
     for key in ['webpage_url', 'extractor', 'extractor_key']: