lbry-sdk/tests/integration/test_streaming.py
Jack Robison a2a1ef34ab
fix tests
2019-05-08 14:55:09 -04:00

367 lines
17 KiB
Python

import os
import hashlib
import aiohttp
import aiohttp.web
from lbrynet.utils import aiohttp_request
from lbrynet.blob.blob_file import MAX_BLOB_SIZE
from lbrynet.testcase import CommandTestCase
def get_random_bytes(n: int) -> bytes:
result = b''.join(hashlib.sha256(os.urandom(4)).digest() for _ in range(n // 16))
if len(result) < n:
result += os.urandom(n - len(result))
elif len(result) > n:
result = result[:-(len(result) - n)]
assert len(result) == n, (n, len(result))
return result
class RangeRequests(CommandTestCase):
async def _restart_stream_manager(self):
self.daemon.stream_manager.stop()
await self.daemon.stream_manager.start()
return
async def _setup_stream(self, data: bytes, save_blobs: bool = True, save_files: bool = False, file_size=0):
self.daemon.conf.save_blobs = save_blobs
self.daemon.conf.save_files = save_files
self.data = data
await self.stream_create('foo', '0.01', data=self.data, file_size=file_size)
if save_blobs:
self.assertTrue(len(os.listdir(self.daemon.blob_manager.blob_dir)) > 1)
await self.daemon.jsonrpc_file_list()[0].fully_reflected.wait()
await self.daemon.jsonrpc_file_delete(delete_from_download_dir=True, claim_name='foo')
self.assertEqual(0, len(os.listdir(self.daemon.blob_manager.blob_dir)))
# await self._restart_stream_manager()
await self.daemon.runner.setup()
site = aiohttp.web.TCPSite(self.daemon.runner, self.daemon.conf.api_host, self.daemon.conf.api_port)
await site.start()
self.assertListEqual(self.daemon.jsonrpc_file_list(), [])
async def _test_range_requests(self):
name = 'foo'
url = f'http://{self.daemon.conf.api_host}:{self.daemon.conf.api_port}/get/{name}'
async with aiohttp_request('get', url) as req:
self.assertEqual(req.headers.get('Content-Type'), 'application/octet-stream')
content_range = req.headers.get('Content-Range')
content_length = int(req.headers.get('Content-Length'))
streamed_bytes = await req.content.read()
self.assertEqual(content_length, len(streamed_bytes))
return streamed_bytes, content_range, content_length
async def test_range_requests_2_byte(self):
self.data = b'hi'
await self._setup_stream(self.data)
streamed, content_range, content_length = await self._test_range_requests()
self.assertEqual(15, content_length)
self.assertEqual(b'hi\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', streamed)
self.assertEqual('bytes 0-14/15', content_range)
async def test_range_requests_15_byte(self):
self.data = b'123456789abcdef'
await self._setup_stream(self.data)
streamed, content_range, content_length = await self._test_range_requests()
self.assertEqual(15, content_length)
self.assertEqual(15, len(streamed))
self.assertEqual(self.data, streamed)
self.assertEqual('bytes 0-14/15', content_range)
async def test_range_requests_0_padded_bytes(self, size: int = (MAX_BLOB_SIZE - 1) * 4,
expected_range: str = 'bytes 0-8388603/8388604', padding=b'',
file_size=0):
self.data = get_random_bytes(size)
await self._setup_stream(self.data, file_size=file_size)
streamed, content_range, content_length = await self._test_range_requests()
self.assertEqual(len(self.data + padding), content_length)
self.assertEqual(streamed, self.data + padding)
self.assertEqual(expected_range, content_range)
async def test_range_requests_1_padded_bytes(self):
await self.test_range_requests_0_padded_bytes(
((MAX_BLOB_SIZE - 1) * 4) - 1, padding=b'\x00'
)
async def test_range_requests_2_padded_bytes(self):
await self.test_range_requests_0_padded_bytes(
((MAX_BLOB_SIZE - 1) * 4) - 2, padding=b'\x00' * 2
)
async def test_range_requests_14_padded_bytes(self):
await self.test_range_requests_0_padded_bytes(
((MAX_BLOB_SIZE - 1) * 4) - 14, padding=b'\x00' * 14
)
async def test_range_requests_no_padding_size_from_claim(self):
size = ((MAX_BLOB_SIZE - 1) * 4) - 14
await self.test_range_requests_0_padded_bytes(size, padding=b'', file_size=size,
expected_range=f"bytes 0-{size-1}/{size}")
async def test_range_requests_15_padded_bytes(self):
await self.test_range_requests_0_padded_bytes(
((MAX_BLOB_SIZE - 1) * 4) - 15, padding=b'\x00' * 15
)
async def test_range_requests_last_block_of_last_blob_padding(self):
self.data = get_random_bytes(((MAX_BLOB_SIZE - 1) * 4) - 16)
await self._setup_stream(self.data)
streamed, content_range, content_length = await self._test_range_requests()
self.assertEqual(len(self.data), content_length)
self.assertEqual(streamed, self.data)
self.assertEqual('bytes 0-8388587/8388588', content_range)
async def test_streaming_only_with_blobs(self):
self.data = get_random_bytes((MAX_BLOB_SIZE - 1) * 4)
await self._setup_stream(self.data)
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
# test that repeated range requests do not create duplicate files
for _ in range(3):
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
current_files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
# test that a range request after restart does not create a duplicate file
await self._restart_stream_manager()
current_files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
stream = self.daemon.jsonrpc_file_list()[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
current_files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
async def test_streaming_only_without_blobs(self):
self.data = get_random_bytes((MAX_BLOB_SIZE - 1) * 4)
await self._setup_stream(self.data, save_blobs=False)
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
# test that repeated range requests do not create duplicate files
for _ in range(3):
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
current_files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
# test that a range request after restart does not create a duplicate file
await self._restart_stream_manager()
current_files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNone(stream.download_directory)
self.assertIsNone(stream.full_path)
current_files_in_download_dir = list(os.scandir(os.path.dirname(self.daemon.conf.data_dir)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
async def test_stream_and_save_file_with_blobs(self):
self.data = get_random_bytes((MAX_BLOB_SIZE - 1) * 4)
await self._setup_stream(self.data, save_files=True)
await self._test_range_requests()
streams = self.daemon.jsonrpc_file_list()
self.assertEqual(1, len(streams))
stream = streams[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
full_path = stream.full_path
files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
for _ in range(3):
await self._test_range_requests()
streams = self.daemon.jsonrpc_file_list()
self.assertEqual(1, len(streams))
stream = streams[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
current_files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
await self._restart_stream_manager()
current_files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
streams = self.daemon.jsonrpc_file_list()
self.assertEqual(1, len(streams))
stream = streams[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
await self._test_range_requests()
streams = self.daemon.jsonrpc_file_list()
self.assertEqual(1, len(streams))
stream = streams[0]
self.assertTrue(os.path.isfile(self.daemon.blob_manager.get_blob(stream.sd_hash).file_path))
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
current_files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
with open(stream.full_path, 'rb') as f:
self.assertEqual(self.data, f.read())
async def test_stream_and_save_file_without_blobs(self):
self.data = get_random_bytes((MAX_BLOB_SIZE - 1) * 4)
await self._setup_stream(self.data, save_files=True)
self.daemon.conf.save_blobs = False
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
full_path = stream.full_path
files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
for _ in range(3):
await self._test_range_requests()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
current_files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
await self._restart_stream_manager()
current_files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
streams = self.daemon.jsonrpc_file_list()
self.assertEqual(1, len(streams))
stream = streams[0]
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
await self._test_range_requests()
streams = self.daemon.jsonrpc_file_list()
self.assertEqual(1, len(streams))
stream = streams[0]
self.assertTrue(os.path.isdir(stream.download_directory))
self.assertTrue(os.path.isfile(stream.full_path))
current_files_in_download_dir = list(os.scandir(os.path.dirname(full_path)))
self.assertEqual(
len(files_in_download_dir), len(current_files_in_download_dir)
)
with open(stream.full_path, 'rb') as f:
self.assertEqual(self.data, f.read())
async def test_switch_save_blobs_while_running(self):
await self.test_streaming_only_without_blobs()
self.daemon.conf.save_blobs = True
blobs_in_stream = self.daemon.jsonrpc_file_list()[0].blobs_in_stream
sd_hash = self.daemon.jsonrpc_file_list()[0].sd_hash
start_file_count = len(os.listdir(self.daemon.blob_manager.blob_dir))
await self._test_range_requests()
self.assertEqual(start_file_count + blobs_in_stream, len(os.listdir(self.daemon.blob_manager.blob_dir)))
self.assertEqual(0, self.daemon.jsonrpc_file_list()[0].blobs_remaining)
# switch back
self.daemon.conf.save_blobs = False
await self._test_range_requests()
self.assertEqual(start_file_count + blobs_in_stream, len(os.listdir(self.daemon.blob_manager.blob_dir)))
self.assertEqual(0, self.daemon.jsonrpc_file_list()[0].blobs_remaining)
await self.daemon.jsonrpc_file_delete(delete_from_download_dir=True, sd_hash=sd_hash)
self.assertEqual(start_file_count, len(os.listdir(self.daemon.blob_manager.blob_dir)))
await self._test_range_requests()
self.assertEqual(start_file_count, len(os.listdir(self.daemon.blob_manager.blob_dir)))
self.assertEqual(blobs_in_stream, self.daemon.jsonrpc_file_list()[0].blobs_remaining)
async def test_file_save_streaming_only_save_blobs(self):
await self.test_streaming_only_with_blobs()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNone(stream.full_path)
self.server.stop_server()
await self.daemon.jsonrpc_file_save('test', self.daemon.conf.data_dir)
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNotNone(stream.full_path)
await stream.finished_writing.wait()
with open(stream.full_path, 'rb') as f:
self.assertEqual(self.data, f.read())
await self.daemon.jsonrpc_file_delete(delete_from_download_dir=True, sd_hash=stream.sd_hash)
async def test_file_save_stop_before_finished_streaming_only(self, wait_for_start_writing=False):
await self.test_streaming_only_with_blobs()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNone(stream.full_path)
self.server.stop_server()
await self.daemon.jsonrpc_file_save('test', self.daemon.conf.data_dir)
stream = self.daemon.jsonrpc_file_list()[0]
path = stream.full_path
self.assertIsNotNone(path)
if wait_for_start_writing:
await stream.started_writing.wait()
self.assertTrue(os.path.isfile(path))
await self._restart_stream_manager()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNotNone(stream.full_path)
self.assertFalse(os.path.isfile(path))
if wait_for_start_writing:
await stream.started_writing.wait()
self.assertTrue(os.path.isfile(path))
async def test_file_save_stop_before_finished_streaming_only_wait_for_start(self):
return await self.test_file_save_stop_before_finished_streaming_only(wait_for_start_writing=True)
async def test_file_save_streaming_only_dont_save_blobs(self):
await self.test_streaming_only_without_blobs()
stream = self.daemon.jsonrpc_file_list()[0]
self.assertIsNone(stream.full_path)
await self.daemon.jsonrpc_file_save('test', self.daemon.conf.data_dir)
stream = self.daemon.jsonrpc_file_list()[0]
await stream.finished_writing.wait()
with open(stream.full_path, 'rb') as f:
self.assertEqual(self.data, f.read())