code: def mock_data(
      num_examples: int = 1,
      num_sub_examples: int = 1,
      max_value: Optional[int] = None,
      *,
      policy: MockPolicy = MockPolicy.AUTO,
      as_dataset_fn: Optional[Callable[..., tf.data.Dataset]] = None,
      as_data_source_fn: Optional[Callable[..., Sequence[Any]]] = None,
      data_dir: Optional[str] = None,
      mock_array_record_data_source: Optional[
          PickableDataSourceMock
      ] = None,
  ) -> Iterator[None]:
      """Mock tfds to generate random data.
  
      ### Usage
  
      * Usage (automated):
  
      ```py
      with tfds.testing.mock_data(num_examples=5):
        ds = tfds.load('some_dataset', split='train')
  
        for ex in ds:  # ds will yield randomly generated examples.
          ex
      ```
  
      All calls to `tfds.load`/`tfds.data_source` within the context manager then
      return deterministic mocked data.
  
      * Usage (manual):
  
      For more control over the generated examples, you can
      manually overwrite the `DatasetBuilder._as_dataset` method:
  
      ```py
      def as_dataset(self, *args, **kwargs):
        return tf.data.Dataset.from_generator(
            lambda: ({
                'image': np.ones(shape=(28, 28, 1), dtype=np.uint8),
                'label': i % 10,
            } for i in range(num_examples)),
            output_types=self.info.features.dtype,
            output_shapes=self.info.features.shape,
        )
  
      with mock_data(as_dataset_fn=as_dataset):
        ds = tfds.load('some_dataset', split='train')
  
        for ex in ds:  # ds will yield the fake data example of 'as_dataset'.
          ex
      ```
  
      Similarly, `as_data_source_fn` allows to manually mock `tfds.data_source`. See
      [the
      tests](https://github.com/tensorflow/datasets/tree/master/tensorflow_datasets/testing/mocking_test.py;l=399;rcl=584915359)
      for more examples.
  
      ### Policy
  
      For improved results, you can copy the true metadata files
      (`dataset_info.json`, `label.txt`, vocabulary files) in
      `data_dir/dataset_name/version`. This will allow the mocked dataset to use
      the true metadata computed during generation (split names,...).
  
      If metadata files are not found, then info from the original class will be
      used, but the features computed during generation won't be available (e.g.
      unknown split names, so any splits are accepted).
  
      ### Miscellaneous
  
      * The examples are deterministically generated. Train and test split will
        yield the same examples.
      * The actual examples will be randomly generated using
        `builder.info.features.get_tensor_info()`.
      * Download and prepare step will always be a no-op.
      * Warning: `info.split['train'].num_examples` won't match
        `len(list(ds_train))`
  
      Some of those points could be improved. If you have suggestions, issues with
      this functions, please open a new issue on our Github.
  
      Args:
        num_examples: Number of fake example to generate.
        num_sub_examples: Number of examples to generate in nested Dataset features.
        max_value: The maximum value present in generated tensors; if max_value is
          None or it is set to 0, then random numbers are generated from the range
          from 0 to 255.
        policy: Strategy to use to generate the fake examples. See
          `tfds.testing.MockPolicy`.
        as_dataset_fn: If provided, will replace the default random example
          generator. This function mock the `FileAdapterBuilder._as_dataset`
        as_data_source_fn: If provided, will replace the default random example
          generator for `as_data_source`. The function must return a custom data
          source that implements `__len__` and `__getitem__` (e.g., a list).
        data_dir: Folder containing the metadata file (searched in
          `data_dir/dataset_name/version`). Overwrite `data_dir` kwargs from
          `tfds.load`. Used in `MockPolicy.USE_FILES` mode.
        mock_array_record_data_source: Overwrite a mock for the underlying
          ArrayRecord data source if it is used. Note: If used the same mock will be
          used for all data sources loaded within this context.
  
      Yields:
        None
      """
  
      original_init_fn = dataset_builder.DatasetBuilder.__init__
      original_as_dataset_fn = dataset_builder.DatasetBuilder.as_dataset
      original_builder_from_files = read_only_builder.builder_from_files
      if mock_array_record_data_source and not isinstance(
          mock_array_record_data_source, PickableDataSourceMock
      ):
          raise ValueError(
              "`mock_array_record_data_source` must be a"
              " `tfds.testing.PickableDataSourceMock` because it must be pickable."
          )
  
      def mock_download_and_prepare(self, *args, **kwargs):
          """`builder.download_and_prepare` is a no-op."""
          del self, args, kwargs  # Unused
  
      def _check_policy(self) -> MockPolicy:
          """Checks the policy and returns the actual used policy for mocking."""
          # When `USE_FILES` is used, make sure the metadata actually exists.
          if epath.Path(self.data_dir).exists():
              logging.info(
                  "Metadata found for %s at %s",
                  self.name,
                  self.data_dir,
              )
          else:
              if policy == MockPolicy.USE_FILES:
                  raise ValueError(
                      "TFDS has been mocked with `MockPolicy.USE_FILES`, but metadata "
                      f"files were not found in {self.data_dir}. "
                      "You should copy the real metadata files, so that the dataset "
                      "can be loaded properly, or set the data_dir kwarg of "
                      "tfds.testing.mock_tfds(data_dir=...).\n"
                  )
              if policy == MockPolicy.AUTO:
                  logging.info(
                      "Metadata NOT found for %s at %s. Will use `MockPolicy.USE_CODE.`",
                      self.name,
                      self.data_dir,
                  )
                  return MockPolicy.USE_CODE
          return policy
  
      def mock_as_dataset_base(self, **kwargs):
          """Mocks `builder.as_dataset`."""
          _check_policy(self)
  
          # Info is already restored at this point, so can mock the file system
          # safely in case of `USE_CODE` mode.
          # The only gfile access is to check `self.data_dir` existance.
          with test_utils.MockFs() as fs:
              fs.add_file(os.path.join(self.data_dir, "tmp.txt"))
              return original_as_dataset_fn(self, **kwargs)
  
      def mock_as_dataset(
          self, split, decoders=None, read_config=None, **kwargs
      ):
          """Mocks `builder._as_dataset`."""
          del kwargs
  
          (
              generator_cls,
              features,
              specs,
              decode_fn,
          ) = _get_fake_data_components(decoders, self.info.features)
  
          ds = tf.data.Dataset.from_generator(
              # `from_generator` takes a callable with signature () -> iterable
              # Recreating a new generator each time ensure that all pipelines are
              # using the same examples
              functools.partial(
                  generator_cls,
                  features=features,
                  num_examples=num_examples,
                  num_sub_examples=num_sub_examples,
                  max_value=max_value,
              ),
              # pylint: enable=g-long-lambda]
              output_types=tf.nest.map_structure(
                  lambda t: t.dtype, specs
              ),
              output_shapes=tf.nest.map_structure(
                  lambda t: t.shape, specs
              ),
          )
          ds = ds.apply(
              tf.data.experimental.assert_cardinality(num_examples)
          )
          ds = ds.map(
              decode_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
  
          if read_config and read_config.add_tfds_id:
              ds_id = reader_lib._make_id_dataset(  # pylint: disable=protected-access
                  filename=f"{self.name}-{split}.tfrecord-00000-of-00001",
                  start_index=0,  # pytype: disable=wrong-arg-types
              )
              ds = tf.data.Dataset.zip((ds, ds_id))
              ds = ds.map(lambda ex, id: {"tfds_id": id, **ex})
  
          return ds
  
      def mock_as_data_source(self, split, decoders=None, **kwargs):
          """Mocks `builder.as_data_source`."""
          del kwargs
          nonlocal mock_array_record_data_source
          mock_data_source = (
              mock_array_record_data_source or PickableDataSourceMock()
          )
          actual_policy = _check_policy(self)
          if split is None:
              split = {s: s for s in self.info.splits}
  
          generator_cls, features, _, _ = _get_fake_data_components(
              decoders, self.info.features
          )
          generator = generator_cls(features, num_examples)
  
          if actual_policy == MockPolicy.USE_CODE:
              # Mock `SplitDict`, because splits are not known in advance.
              split_dict = mock.MagicMock(spec=splits_lib.SplitDict)
          else:
              # Use the real `SplitDict`.
              split_dict = splits_lib.SplitDict
          with mock.patch.object(
              splits_lib,
              "SplitDict",
              new=split_dict,
          ), mock.patch(
              "array_record.python.array_record_data_source.ArrayRecordDataSource",
              mock_data_source,
          ), mock.patch(
              "tensorflow_datasets.core.dataset_info.DatasetInfo.file_format",
              new_callable=mock.PropertyMock,
              # Force ARRAY_RECORD as the default file_format.
              return_value=file_adapters.FileFormat.ARRAY_RECORD,
          ):
              self.info.features.deserialize_example_np = (
                  _deserialize_example_np
              )
              mock_data_source.return_value.__len__.return_value = (
                  num_examples
              )
              mock_data_source.return_value._generator = (  # pylint:disable=protected-access
                  generator
              )
              mock_data_source.return_value.__getitem__ = (
                  functools.partial(_getitem, generator=generator)
              )
              mock_data_source.return_value.__getitems__ = (
                  functools.partial(_getitems, generator=generator)
              )
  
              def build_single_data_source(split):
                  single_data_source = (
                      array_record.ArrayRecordDataSource(
                          dataset_info=self.info,
                          split=split,
                          decoders=decoders,
                      )
                  )
                  return single_data_source
  
              return tree_utils.map_structure(
                  build_single_data_source, split
              )
  
      if not as_dataset_fn:
          as_dataset_fn = mock_as_dataset
  
      if policy == MockPolicy.USE_CODE:
          mock_data_dir = "/tmp/non-existing/tfds/data_dir/"
      elif not data_dir:
          mock_data_dir = os.path.join(
              os.path.dirname(__file__), "metadata"
          )
      else:  # AUTO or USE_FILES with explicitly given `data_dir`
          mock_data_dir = data_dir
  
      def mock_init(*args, data_dir=None, **kwargs):
          del data_dir  # Unused. Inject `mock_data_dir` instead.
          return original_init_fn(
              *args, data_dir=mock_data_dir, **kwargs
          )
  
      def new_builder_from_files(*args, **kwargs):
          # Replace the user-given data dir by the mocked one
          kwargs.pop("data_dir", None)
          # `DatasetBuilder.__init__` is mocked above to inject the wrong data_dir.
          # So we restore the original `DatasetBuilder.__init__` inside
          # `builder_from_files` calls.
          with mock.patch(
              f"{core}.dataset_builder.DatasetBuilder.__init__",
              original_init_fn,
          ):
              return original_builder_from_files(
                  *args, data_dir=mock_data_dir, **kwargs
              )
  
      core = "tensorflow_datasets.core"
      with contextlib.ExitStack() as stack:
          for path, mocked_fn in [
              # All GCS access should fail (is_dataset_on_gcs,...).
              (f"{core}.utils.gcs_utils.exists", lambda path: False),
              # Patch `data_dir`: `data_dir` explicitly set by users will be ignored.
              # `data_dir` is used at two places:
              # * `builder_from_files` to search read-only datasets loaded from config
              # * `DatasetBuilder.__init__` otherwise
              (
                  f"{core}.dataset_builder.DatasetBuilder.__init__",
                  mock_init,
              ),
              (
                  f"{core}.read_only_builder.builder_from_files",
                  new_builder_from_files,
              ),
              # Patch DatasetBuilder
              (
                  f"{core}.dataset_builder.DatasetBuilder.download_and_prepare",
                  mock_download_and_prepare,
              ),
              (
                  f"{core}.dataset_builder.DatasetBuilder.as_dataset",
                  mock_as_dataset_base,
              ),
              (
                  f"{core}.dataset_builder.DatasetBuilder.as_data_source",
                  as_data_source_fn
                  if as_data_source_fn
                  else mock_as_data_source,
              ),
              (
                  f"{core}.dataset_builder.FileReaderBuilder._as_dataset",
                  as_dataset_fn,
              ),
          ]:
              stack.enter_context(mock.patch(path, mocked_fn))
          yield
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import enum
  import functools
  import os
  from typing import Any, Callable, Iterator, Optional
  from unittest import mock
  from absl import logging
  from etils import epath
  import numpy
  from tensorflow_datasets.core import dataset_builder
  from tensorflow_datasets.core import decode
  from tensorflow_datasets.core import features
  from tensorflow_datasets.core import file_adapters
  from tensorflow_datasets.core import read_only_builder
  from tensorflow_datasets.core import reader
  from tensorflow_datasets.core import splits
  from tensorflow_datasets.core.data_sources import array_record
  from tensorflow_datasets.core.utils import dtype_utils
  from tensorflow_datasets.core.utils import tree_utils
  from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow
  from tensorflow_datasets.testing import test_utils
label_desc: Use image node_id 5 for calling a global function with example usage: mock_data(num_examples, num_sub_examples, max_value) and returns: policy, ds, original_init_fn, original_as_dataset_fn, tree_utils, original_builder_from_files, MockPolicy, single_data_source