test/test_downloader_http.py

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