From e862c99f6c5588a02279330d27e6dae0ee543461 Mon Sep 17 00:00:00 2001 From: Victor Shyba Date: Thu, 29 Sep 2022 17:29:26 -0300 Subject: [PATCH] generate 3 files, check that streamed is the largest, add method to list files --- lbry/torrent/session.py | 20 ++++++++++++++----- .../datanetwork/test_file_commands.py | 2 ++ 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/lbry/torrent/session.py b/lbry/torrent/session.py index a145a3132..a8308d449 100644 --- a/lbry/torrent/session.py +++ b/lbry/torrent/session.py @@ -4,7 +4,7 @@ import os import logging import random from tempfile import mkdtemp -from typing import Optional, Tuple +from typing import Optional, Tuple, Dict import libtorrent @@ -154,7 +154,7 @@ class TorrentSession: self.tasks = [] self.wait_start = True - async def add_fake_torrent(self, file_count=1): + async def add_fake_torrent(self, file_count=3): tmpdir = mkdtemp() info = _create_fake_torrent(tmpdir, file_count=file_count) flags = libtorrent.add_torrent_params_flags_t.flag_seed_mode @@ -259,20 +259,30 @@ class TorrentSession: handle = self._handles[btih] return handle.stream_range_as_completed(handle.largest_file_index, start, end) + def get_files(self, btih) -> Dict: + handle = self._handles[btih] + return { + handle.torrent_file.file_path(file_num): handle.torrent_file.file_size(file_num) + for file_num in range(handle.torrent_file.num_files()) + if '.pad' not in handle.torrent_file.file_path(file_num) + } + def get_magnet_uri(btih): return f"magnet:?xt=urn:btih:{btih}" -def _create_fake_torrent(tmpdir, file_count=1): - # layout: subdir/tmp{0..file_count-1} 40k files. v1+v2. automatic piece size. +def _create_fake_torrent(tmpdir, file_count=3, largest_index=1): + # layout: subdir/tmp{0..file_count-1} files. v1+v2. automatic piece size. + # largest_index: which file index {0 ... file_count} will be the largest file file_storage = libtorrent.file_storage() subfolder = os.path.join(tmpdir, "subdir") os.mkdir(subfolder) for file_number in range(file_count): file_name = f"tmp{file_number}" with open(os.path.join(subfolder, file_name), 'wb') as myfile: - size = myfile.write(bytes([random.randint(0, 255) for _ in range(40)]) * 1024) + size = myfile.write( + bytes([random.randint(0, 255) for _ in range(10 - abs(file_number - largest_index))]) * 1024) file_storage.add_file(os.path.join("subdir", file_name), size) t = libtorrent.create_torrent(file_storage, 0, 0) libtorrent.set_piece_hashes(t, tmpdir) diff --git a/tests/integration/datanetwork/test_file_commands.py b/tests/integration/datanetwork/test_file_commands.py index d3af87fc0..fdcdff355 100644 --- a/tests/integration/datanetwork/test_file_commands.py +++ b/tests/integration/datanetwork/test_file_commands.py @@ -71,6 +71,8 @@ class FileCommands(CommandTestCase): self.assertEqual(content_length, len(streamed_bytes)) self.assertEqual(f"bytes 0-{expected_size - 1}/{expected_size}", content_range) + self.assertEqual(len(streamed_bytes), max(self.seeder_session.get_files(btih).values())) + @skipIf(TorrentSession is None, "libtorrent not installed") async def test_download_torrent(self): tx, btih = await self.initialize_torrent()