test/test_downloader_http.py

   1 #!/usr/bin/env python3
   2 # Allow direct execution
   3 import os
   4 import re
   5 import sys
   6 import unittest
   7
   8 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
   9
  10 import threading
  11 from test.helper import http_server_port, try_rm
  12
  13 from yt_dlp import YoutubeDL
  14 from yt_dlp.compat import compat_http_server
  15 from yt_dlp.downloader.http import HttpFD
  16 from yt_dlp.utils import encodeFilename
  17
  18 TEST_DIR = os.path.dirname(os.path.abspath(__file__))
  19
  20
  21 TEST_SIZE = 10 * 1024
  22
  23
  24 class HTTPTestRequestHandler(compat_http_server.BaseHTTPRequestHandler):
  25     def log_message(self, format, *args):
  26         pass
  27
  28     def send_content_range(self, total=None):
  29         range_header = self.headers.get('Range')
  30         start = end = None
  31         if range_header:
  32             mobj = re.search(r'^bytes=(\d+)-(\d+)', range_header)
  33             if mobj:
  34                 start = int(mobj.group(1))
  35                 end = int(mobj.group(2))
  36         valid_range = start is not None and end is not None
  37         if valid_range:
  38             content_range = 'bytes %d-%d' % (start, end)
  39             if total:
  40                 content_range += '/%d' % total
  41             self.send_header('Content-Range', content_range)
  42         return (end - start + 1) if valid_range else total
  43
  44     def serve(self, range=True, content_length=True):
  45         self.send_response(200)
  46         self.send_header('Content-Type', 'video/mp4')
  47         size = TEST_SIZE
  48         if range:
  49             size = self.send_content_range(TEST_SIZE)
  50         if content_length:
  51             self.send_header('Content-Length', size)
  52         self.end_headers()
  53         self.wfile.write(b'#' * size)
  54
  55     def do_GET(self):
  56         if self.path == '/regular':
  57             self.serve()
  58         elif self.path == '/no-content-length':
  59             self.serve(content_length=False)
  60         elif self.path == '/no-range':
  61             self.serve(range=False)
  62         elif self.path == '/no-range-no-content-length':
  63             self.serve(range=False, content_length=False)
  64         else:
  65             assert False
  66
  67
  68 class FakeLogger:
  69     def debug(self, msg):
  70         pass
  71
  72     def warning(self, msg):
  73         pass
  74
  75     def error(self, msg):
  76         pass
  77
  78
  79 class TestHttpFD(unittest.TestCase):
  80     def setUp(self):
  81         self.httpd = compat_http_server.HTTPServer(
  82             ('127.0.0.1', 0), HTTPTestRequestHandler)
  83         self.port = http_server_port(self.httpd)
  84         self.server_thread = threading.Thread(target=self.httpd.serve_forever)
  85         self.server_thread.daemon = True
  86         self.server_thread.start()
  87
  88     def download(self, params, ep):
  89         params['logger'] = FakeLogger()
  90         ydl = YoutubeDL(params)
  91         downloader = HttpFD(ydl, params)
  92         filename = 'testfile.mp4'
  93         try_rm(encodeFilename(filename))
  94         self.assertTrue(downloader.real_download(filename, {
  95             'url': 'http://127.0.0.1:%d/%s' % (self.port, ep),
  96         }))
  97         self.assertEqual(os.path.getsize(encodeFilename(filename)), TEST_SIZE)
  98         try_rm(encodeFilename(filename))
  99
 100     def download_all(self, params):
 101         for ep in ('regular', 'no-content-length', 'no-range', 'no-range-no-content-length'):
 102             self.download(params, ep)
 103
 104     def test_regular(self):
 105         self.download_all({})
 106
 107     def test_chunked(self):
 108         self.download_all({
 109             'http_chunk_size': 1000,
 110         })
 111
 112
 113 if __name__ == '__main__':
 114     unittest.main()