code: def testSaveRestoreOneShotIterator(self, enable_async_ckpt):
      checkpoint_directory = self.get_temp_dir()
      checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
      dataset = (
          dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6])
          .map(math_ops.square)
          .batch(2)
      )
      iterator = iter(dataset)
      get_next = iterator.get_next
      ckpt_options = checkpoint_options.CheckpointOptions(
          experimental_enable_async_checkpoint=enable_async_ckpt
      )
      checkpoint = trackable_utils.Checkpoint(iterator=iterator)
      self.assertAllEqual([1, 4], get_next())
      save_path = checkpoint.save(
          checkpoint_prefix, options=ckpt_options
      )
      self.assertAllEqual([9, 16], get_next())
      self.assertAllEqual([25, 36], get_next())
      checkpoint.restore(save_path).run_restore_ops()
      self.assertAllEqual([9, 16], get_next())
      self.assertAllEqual([25, 36], get_next())
      with self.assertRaises(errors.OutOfRangeError):
          get_next()
  
imports: import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.checkpoint import checkpoint_options
  from tensorflow.python.data.experimental.ops import grouping
  from tensorflow.python.data.experimental.ops import interleave_ops
  from tensorflow.python.data.experimental.ops import scan_ops
  from tensorflow.python.data.experimental.ops import take_while_ops
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import options
  from tensorflow.python.eager import test
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import gen_dataset_ops
  from tensorflow.python.ops import io_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import parsing_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import script_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import gfile
label_desc: Use image node_id 10 for calling the CheckpointTest obj's underlying member method code with example usage: obj.testSaveRestoreOneShotIterator(enable_async_ckpt) without return types