class_docstr: Tests for KafkaDataset.
docstr: Tests for KafkaDataset save and restore.
docstr_tok: ["Tests","for","KafkaDataset","save","and","restore","."]
code: def test_kafka_dataset_save_and_restore(self):
      """Tests for KafkaDataset save and restore."""
      tf.compat.v1.disable_eager_execution()
  
      import tensorflow_io.kafka as kafka_io  # pylint: disable=wrong-import-position
  
      g = tf.Graph()
      with g.as_default():
          topics = tf.compat.v1.placeholder(
              tf.dtypes.string, shape=[None]
          )
          num_epochs = tf.compat.v1.placeholder(
              tf.dtypes.int64, shape=[]
          )
  
          repeat_dataset = kafka_io.KafkaDataset(
              topics, group="test", eof=True
          ).repeat(num_epochs)
          iterator = repeat_dataset.make_initializable_iterator()
          get_next = iterator.get_next()
  
          it = tf.data.experimental.make_saveable_from_iterator(
              iterator
          )
          g.add_to_collection(
              tf.compat.v1.GraphKeys.SAVEABLE_OBJECTS, it
          )
          saver = tf.compat.v1.train.Saver()
  
          model_file = "/tmp/test-kafka-model"
          with self.cached_session() as sess:
              sess.run(
                  iterator.initializer,
                  feed_dict={topics: ["test:0:0:4"], num_epochs: 1},
              )
              for i in range(3):
                  self.assertEqual(
                      ("D" + str(i)).encode(), sess.run(get_next)
                  )
              # Save current offset which is 2
              saver.save(sess, model_file, global_step=3)
  
          checkpoint_file = "/tmp/test-kafka-model-3"
          with self.cached_session() as sess:
              saver.restore(sess, checkpoint_file)
              # Restore current offset to 2
              for i in [2, 3]:
                  self.assertEqual(
                      ("D" + str(i)).encode(), sess.run(get_next)
                  )
  
file_path: io/tests/test_kafka_v1.py
imports: import time
  import pytest
  import tensorflow
label_desc: Use image node_id 2 for calling the KafkaDatasetTest obj's underlying member method code with example usage: obj.test_kafka_dataset_save_and_restore() without return types