Merge pull request #997 from reef-technologies/fix_upload_threads

fix upload threads setting

b2/_internal/console_tool.py

@@ -734,9 +734,8 @@ def _get_threads_from_args(self, args) -> int:
 
  def _set_threads_from_args(self, args):
  threads = self._get_threads_from_args(args)
- # FIXME: This is using deprecated API. It should be be replaced when moving to b2sdk apiver 3.
- # There is `max_download_workers` param in B2Api constructor for this.
  self.api.services.download_manager.set_thread_pool_size(threads)
+ self.api.services.upload_manager.set_thread_pool_size(threads)
 
 
 class _TqdmCloser:
@@ -2811,8 +2810,6 @@ def _run(self, args):
  upload_threads = args.uploadThreads
  download_threads = args.downloadThreads
 
- # FIXME: This is using deprecated API. It should be replaced when moving to b2sdk apiver 3.
- # There are `max_X_workers` params in B2Api constructor for this.
  self.api.services.upload_manager.set_thread_pool_size(upload_threads)
  self.api.services.download_manager.set_thread_pool_size(download_threads)
 

changelog.d/+upload_threads.fixed.md

@@ -0,0 +1 @@
+Fix `--threads` option being silently ignored in upload commands.

pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
 dependencies = [
  "argcomplete>=2,<4",
  "arrow>=1.0.2,<2.0.0",
- "b2sdk>=1.31.0,<2",
+ "b2sdk>=1.32.0,<2",
  "docutils>=0.18.1",
  "idna~=3.4; platform_system == 'Java'",
  "importlib-metadata>=3.3; python_version < '3.8'",

test/unit/console_tool/test_download_file.py

@@ -190,3 +190,19 @@ def test_cat__b2id_uri(b2_cli, bucket, uploaded_stdout_txt, tmp_path, capfd):
  """Test download_file_by_name stdout alias support"""
  b2_cli.run(['cat', '--noProgress', "b2id://9999"],)
  assert capfd.readouterr().out == uploaded_stdout_txt['content']
+
+
+def test__download_file__threads(b2_cli, local_file, uploaded_file, tmp_path):
+ num_threads = 13
+ output_path = tmp_path / 'output.txt'
+
+ b2_cli.run(
+ [
+ 'download-file', '--noProgress', '--threads',
+ str(num_threads), 'b2://my-bucket/file1.txt',
+ str(output_path)
+ ]
+ )
+
+ assert output_path.read_text() == uploaded_file['content']
+ assert b2_cli.b2_api.services.download_manager.get_thread_pool_size() == num_threads

test/unit/console_tool/test_upload_file.py

@@ -125,3 +125,34 @@ def test_upload_file__stdin(b2_cli, bucket, tmpdir, mock_stdin):
  remove_version=True,
  expected_part_of_stdout=expected_stdout,
  )
+
+
+def test_upload_file__threads_setting(b2_cli, bucket, tmp_path):
+ """Test upload_file supports setting number of threads"""
+ num_threads = 66
+ filename = 'file1.txt'
+ content = 'hello world'
+ local_file1 = tmp_path / 'file1.txt'
+ local_file1.write_text(content)
+
+ expected_json = {
+ "action": "upload",
+ "contentSha1": "2aae6c35c94fcfb415dbe95f408b9ce91ee846ed",
+ "fileInfo": {
+ "src_last_modified_millis": f"{local_file1.stat().st_mtime_ns // 1000000}"
+ },
+ "fileName": filename,
+ "size": len(content),
+ }
+
+ b2_cli.run(
+ [
+ 'upload-file', '--noProgress', 'my-bucket', '--threads',
+ str(num_threads),
+ str(local_file1), 'file1.txt'
+ ],
+ expected_json_in_stdout=expected_json,
+ remove_version=True,
+ )
+
+ assert b2_cli.b2_api.services.upload_manager.get_thread_pool_size() == num_threads

test/unit/test_console_tool.py

@@ -242,7 +242,7 @@ def _run_command(
  self.assertNotIn(unexpected_part_of_stdout, actual_stdout)
  if expected_stderr is not None:
  self.assertEqual(expected_stderr, actual_stderr, 'stderr')
- self.assertEqual(expected_status, actual_status, 'exit status code')
+ assert expected_status == actual_status, 'exit status code'
  return actual_status, actual_stdout, actual_stderr
 
  @classmethod
@@ -1094,43 +1094,6 @@ def test_files_encrypted(self):
  expected_json_in_stdout=expected_json,
  )
 
- def _test_download_threads(self, download_by, num_threads):
- self._authorize_account()
- self._create_my_bucket()
-
- with TempDir() as temp_dir:
- local_file = self._make_local_file(temp_dir, 'file.txt')
- self._run_command(
- ['upload-file', '--noProgress', 'my-bucket', local_file, 'file.txt'],
- remove_version=True,
- )
-
- command = [
- 'download-file-by-%s' % download_by, '--noProgress', '--threads',
- str(num_threads)
- ]
- command += ['9999'] if download_by == 'id' else ['my-bucket', 'file.txt']
- local_download = os.path.join(temp_dir, 'download.txt')
- command += [local_download]
- self._run_command(
- command,
- expected_stderr=
- f'WARNING: download-file-by-{download_by} command is deprecated. Use download-file instead.\n'
- )
- self.assertEqual(b'hello world', self._read_file(local_download))
-
- def test_download_by_id_1_thread(self):
- self._test_download_threads(download_by='id', num_threads=1)
-
- def test_download_by_id_10_threads(self):
- self._test_download_threads(download_by='id', num_threads=10)
-
- def test_download_by_name_1_thread(self):
- self._test_download_threads(download_by='name', num_threads=1)
-
- def test_download_by_name_10_threads(self):
- self._test_download_threads(download_by='name', num_threads=10)
-
  def _test_download_to_directory(self, download_by: str):
  self._authorize_account()
  self._create_my_bucket()