code: def test_batch_block_refs():
      with mock.patch(
          "ray.data._internal.block_batching.block_batching._prefetch_blocks"
      ) as mock_prefetch, mock.patch(
          "ray.data._internal.block_batching.block_batching.batch_blocks"
      ) as mock_batch_blocks:
          block_iter = block_generator(2, 2)
          batch_iter = batch_block_refs(block_iter)
          for _ in batch_iter:
              pass
          assert mock_prefetch.call_count == 1
          assert mock_batch_blocks.call_count == 1
  
file_path: ray/python/ray/data/tests/block_batching/test_block_batching.py
imports: from typing import List
  from unittest import mock
  import pyarrow
  import pytest
  from ray.data._internal.block_batching.block_batching import BlockPrefetcher, _prefetch_blocks, batch_block_refs, batch_blocks
  from ray.data.block import Block
label_desc: Use image node_id 2 for calling a global function with example usage: test_batch_block_refs() without return types