code: def test_kafka_output_sequence():
      """Test case based on fashion mnist tutorial"""
      fashion_mnist = tf.keras.datasets.fashion_mnist
      (
          (train_images, train_labels),
          (test_images, _),
      ) = fashion_mnist.load_data()
  
      class_names = [
          "T-shirt/top",
          "Trouser",
          "Pullover",
          "Dress",
          "Coat",
          "Sandal",
          "Shirt",
          "Sneaker",
          "Bag",
          "Ankle boot",
      ]
  
      train_images = train_images / 255.0
      test_images = test_images / 255.0
  
      model = tf.keras.Sequential(
          [
              tf.keras.layers.Flatten(input_shape=(28, 28)),
              tf.keras.layers.Dense(128, activation=tf.nn.relu),
              tf.keras.layers.Dense(10, activation=tf.nn.softmax),
          ]
      )
  
      model.compile(
          optimizer="adam",
          loss="sparse_categorical_crossentropy",
          metrics=["accuracy"],
      )
  
      model.fit(train_images, train_labels, epochs=5)
  
      class OutputCallback(tf.keras.callbacks.Callback):
          """KafkaOutputCallback"""
  
          def __init__(
              self, batch_size, topic, servers
          ):  # pylint: disable=super-init-not-called
              self._sequence = kafka_ops.KafkaOutputSequence(
                  topic=topic, servers=servers
              )
              self._batch_size = batch_size
  
          def on_predict_batch_end(self, batch, logs=None):
              index = batch * self._batch_size
              for outputs in logs["outputs"]:
                  for output in outputs:
                      self._sequence.setitem(
                          index, class_names[np.argmax(output)]
                      )
                      index += 1
  
          def flush(self):
              self._sequence.flush()
  
      channel = f"e{time.time()}e"
      topic = "test_" + channel
  
      # By default batch size is 32
      output = OutputCallback(32, topic, "localhost")
      predictions = model.predict(test_images, callbacks=[output])
      output.flush()
  
      predictions = [
          class_names[v] for v in np.argmax(predictions, axis=1)
      ]
  
      # Reading from `test_e(time)e` we should get the same result
      dataset = tfio.kafka.KafkaDataset(
          topics=[topic], group="test", eof=True
      )
      for entry, prediction in zip(dataset, predictions):
          assert entry.numpy() == prediction.encode()
  
imports: import sys
  import time
  import pytest
  import numpy
  import threading
  import tensorflow
  import tensorflow_io
label_desc: Use image node_id 2 for calling a global function with example usage: test_kafka_output_sequence() without return types