test_downloader_http.py 3.4 KB

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