lbry-sdk/tests/integration/test_streaming.py

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):
        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)
        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''):
        self.data = get_random_bytes(size)
        await self._setup_stream(self.data)
        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_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.assertIsNone(stream.full_path)
        self.assertFalse(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())