_ Git - youtube-dl/blob - test/test_downloader_http.py

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