3 files changed, 74 insertions, 27 deletions
diff --git a/devscripts/run_tests.bat b/devscripts/run_tests.bat
index 190d23918..57b1f4bf4 100644
--- a/devscripts/run_tests.bat
+++ b/devscripts/run_tests.bat
@@ -1,17 +1,4 @@
-@setlocal
 @echo off
-cd /d %~dp0..
 
-if ["%~1"]==[""] (
-    set "test_set="test""
-) else if ["%~1"]==["core"] (
-    set "test_set="-m not download""
-) else if ["%~1"]==["download"] (
-    set "test_set="-m "download""
-) else (
-    echo.Invalid test type "%~1". Use "core" ^| "download"
-    exit /b 1
-)
-
-set PYTHONWARNINGS=error
-pytest %test_set%
+>&2 echo run_tests.bat is deprecated. Please use `devscripts/run_tests.py` instead
+python %~dp0run_tests.py %~1
diff --git a/devscripts/run_tests.py b/devscripts/run_tests.py
new file mode 100755
index 000000000..b0c6ee67a
--- /dev/null
+++ b/devscripts/run_tests.py
@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+
+import argparse
+import functools
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+
+fix_test_name = functools.partial(re.compile(r'IE(_all|_\d+)?$').sub, r'\1')
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Run selected yt-dlp tests')
+    parser.add_argument(
+        'test', help='a extractor tests, or one of "core" or "download"', nargs='*')
+    parser.add_argument(
+        '-k', help='run a test matching EXPRESSION. Same as "pytest -k"', metavar='EXPRESSION')
+    return parser.parse_args()
+
+
+def run_tests(*tests, pattern=None):
+    run_core = 'core' in tests or (not pattern and not tests)
+    run_download = 'download' in tests
+    tests = list(map(fix_test_name, tests))
+
+    arguments = ['pytest', '-Werror', '--tb', 'short']
+    if run_core:
+        arguments.extend(['-m', 'not download'])
+    elif run_download:
+        arguments.extend(['-m', 'download'])
+    elif pattern:
+        arguments.extend(['-k', pattern])
+    else:
+        arguments.extend(
+            f'test/test_download.py::TestDownload::test_{test}' for test in tests)
+
+    print(f'Running {arguments}')
+    try:
+        subprocess.run(arguments)
+        return
+    except FileNotFoundError:
+        pass
+
+    arguments = [sys.executable, '-Werror', '-m', 'unittest']
+    if run_core:
+        print('"pytest" needs to be installed to run core tests', file=sys.stderr)
+        return
+    elif run_download:
+        arguments.append('test.test_download')
+    elif pattern:
+        arguments.extend(['-k', pattern])
+    else:
+        arguments.extend(
+            f'test.test_download.TestDownload.test_{test}' for test in tests)
+
+    print(f'Running {arguments}')
+    subprocess.run(arguments)
+
+
+if __name__ == '__main__':
+    try:
+        args = parse_args()
+
+        os.chdir(Path(__file__).parent.parent)
+        run_tests(*args.test, pattern=args.k)
+    except KeyboardInterrupt:
+        pass
diff --git a/devscripts/run_tests.sh b/devscripts/run_tests.sh
index faa642e96..123ceb1ee 100755
--- a/devscripts/run_tests.sh
+++ b/devscripts/run_tests.sh
@@ -1,14 +1,4 @@
 #!/usr/bin/env sh
 
-if [ -z "$1" ]; then
-    test_set='test'
-elif [ "$1" = 'core' ]; then
-    test_set="-m not download"
-elif [ "$1" = 'download' ]; then
-    test_set="-m download"
-else
-    echo 'Invalid test type "'"$1"'". Use "core" | "download"'
-    exit 1
-fi
-
-python3 -bb -Werror -m pytest "$test_set"
+>&2 echo 'run_tests.sh is deprecated. Please use `devscripts/run_tests.py` instead'
+python3 devscripts/run_tests.py "$1"