code: def test_beam_datasets(
      tmp_path: pathlib.Path,
      dataset_cls: dataset_builder.GeneratorBasedBuilder,
      make_dl_config: Callable[[], download.DownloadConfig],
  ):
      dataset_name = dataset_cls.name
  
      builder = dataset_cls(data_dir=tmp_path)
      builder.download_and_prepare(download_config=make_dl_config())
  
      data_path = tmp_path / dataset_name / "1.0.0"
      assert data_path.exists()  # Dataset has been generated
  
      # Check number of shards/generated files
      _test_shards(
          data_path,
          pattern="%s-test.tfrecord-{:05}-of-{:05}" % dataset_name,
          # Liquid sharding is not guaranteed to always use the same number.
          num_shards=builder.info.splits["test"].num_shards,
      )
      _test_shards(
          data_path,
          pattern="%s-train.tfrecord-{:05}-of-{:05}" % dataset_name,
          num_shards=1,
      )
  
      ds = dataset_utils.as_numpy(builder.as_dataset())
  
      def get_id(ex):
          return ex["id"]
  
      _assert_values_equal(
          sorted(list(ds["test"]), key=get_id),
          sorted([_gen_example(i)[1] for i in range(725)], key=get_id),
      )
      _assert_values_equal(
          sorted(list(ds["train"]), key=get_id),
          sorted([_gen_example(i)[1] for i in range(1000)], key=get_id),
      )
  
      assert builder.info.metadata == builder.EXPECTED_METADATA
  
imports: import pathlib
  from typing import Callable
  from unittest import mock
  import apache_beam
  from etils import epath
  import numpy
  import pytest
  import tensorflow
  from tensorflow_datasets.core import dataset_builder
  from tensorflow_datasets.core import dataset_info
  from tensorflow_datasets.core import dataset_utils
  from tensorflow_datasets.core import download
  from tensorflow_datasets.core import features
  from tensorflow_datasets.core import utils
label_desc: Use image node_id 5 for calling a global function with example usage: test_beam_datasets(tmp_path, dataset_cls, make_dl_config) and returns: ex