code: def test_format_batches(batch_format):
      block_iter = block_generator(num_rows=2, num_blocks=2)
      batch_iter = (
          Batch(i, block) for i, block in enumerate(block_iter)
      )
      batch_iter = list(
          format_batches(batch_iter, batch_format=batch_format)
      )
  
      for batch in batch_iter:
          if batch_format == "pandas":
              assert isinstance(batch.data, pd.DataFrame)
          elif batch_format == "arrow":
              assert isinstance(batch.data, pa.Table)
          elif batch_format == "numpy":
              assert isinstance(batch.data, dict)
              assert isinstance(batch.data["foo"], np.ndarray)
  
      assert [batch.batch_idx for batch in batch_iter] == list(
          range(len(batch_iter))
      )
  
imports: import threading
  import time
  import numpy
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data._internal.block_batching.interfaces import Batch
  from ray.data._internal.block_batching.util import _calculate_ref_hits, blocks_to_batches, collate, finalize_batches, format_batches, resolve_block_refs
  from ray.data._internal.util import Queue, make_async_gen
label_desc: Use image node_id 4 for calling a global function with example usage: test_format_batches(batch_format) without return types