code: def test_mocking_rlu_nested_dataset():
      """Test of a nested dataset.
  
      In this test we use the dataset rlu_atari.
      The dataset has the following features:
  
        features=tfds.features.FeaturesDict({
          'episode_id': tf.int64,
          'checkpoint_id': tf.int64,
          'episode_return': tf.float32,
          'steps': tfds.features.Dataset({
              'action': tf.int64,
              'discount': tf.float32,
              'is_first': tf.bool,
              'is_last': tf.bool,
              'is_terminal': tf.bool,
              'observation': tfds.features.Image(shape=(84, 84, 1), dtype=tf.uint8),
              'reward': tf.float32,
          }),
        })
      """
      with tfds.testing.mock_data(
          num_examples=3, policy=tfds.testing.MockPolicy.USE_CODE
      ):
          ds = tfds.load("rlu_atari/Pong_run_1", split="train")
  
          steps = ds.element_spec["steps"]
          assert isinstance(steps, tf.data.DatasetSpec)
          assert steps.element_spec["reward"] == tf.TensorSpec(
              shape=(), dtype=tf.float32
          )
  
          for ex in ds.take(3):
              ds_steps = ex["steps"]
              assert isinstance(ds_steps, tf.data.Dataset)
  
              ds_steps_iter = iter(ds_steps)
              steps_ex = next(ds_steps_iter)
              assert set(steps_ex.keys()) == {
                  "action",
                  "discount",
                  "is_first",
                  "is_last",
                  "is_terminal",
                  "observation",
                  "reward",
              }
              assert steps_ex["observation"].shape == (84, 84, 1)
  
imports: import functools
  import pickle
  import numpy
  import pytest
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 16 for calling a global function with example usage: test_mocking_rlu_nested_dataset() without return types