import tensorflow as tf

class MetaGraphTest(tf.test.TestCase):

  def _TestDir(self, test_name):
    test_dir = os.path.join(self.get_temp_dir(), test_name)
    if os.path.exists(test_dir):
      shutil.rmtree(test_dir)
    gfile.MakeDirs(test_dir)
    return test_dir

  def testAddCollectionDef(self):
    test_dir = self._TestDir("good_collection")
    filename = os.path.join(test_dir, "metafile")
    with self.test_session():
      # Creates a graph.
      v0 = tf.Variable(10.0, name="v0")
      var = tf.Variable(tf.constant(0, dtype=tf.int64))
      count_up_to = var.count_up_to(3)
      input_queue = tf.FIFOQueue(30, tf.float32, shared_name="collection_queue")
      qr = tf.train.QueueRunner(input_queue, [count_up_to])
      tf.initialize_all_variables()
      # Creates a saver.
      save = tf.train.Saver({"v0": v0})
      # Adds a set of collections.
      tf.add_to_collection("int_collection", 3)
      tf.add_to_collection("float_collection", 3.5)
      tf.add_to_collection("string_collection", "hello")
      tf.add_to_collection("variable_collection", v0)
      # Add QueueRunners.
      tf.train.add_queue_runner(qr)
      # Adds user_defined proto in three formats: string, bytes and Any.
