import tensorflow as tf

import tensorflow as tf


_SLEEP_TIME = 1.0


class DynamicBatchingTest(tf.test.TestCase):

  def test_one(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(a, b):
        batch_size = tf.shape(a)[0]
        return a + b, tf.tile([batch_size], [batch_size])

      output = f(tf.constant([[1, 3]]), tf.constant([2]))

      tf.train.start_queue_runners()

      result, batch_size = session.run(output)

      self.assertAllEqual([[3, 5]], result)
      self.assertAllEqual([1], batch_size)

  def test_two(self):
    with self.test_session() as session:
