# Path: tensorflowonspark/TFManager.py
# class TFManager(BaseManager):
#   """Python multiprocessing.Manager for distributed, multi-process communication."""
#   pass
# 
# Path: tensorflowonspark/TFNode.py
# TF_VERSION = pkg_resources.get_distribution('tensorflow').version
#   HADOOP_SCHEMES = ['adl://',
#                     'file://',
#                     'hdfs://',
#                     'oss://',
#                     's3://',
#                     's3a://',
#                     's3n://',
#                     'swift://',
#                     'viewfs://',
#                     'wasb://']
# def hdfs_path(ctx, path):
# def start_cluster_server(ctx, num_gpus=1, rdma=False):
# def next_batch(mgr, batch_size, qname='input'):
# def export_saved_model(sess, export_dir, tag_set, signatures):
# def release_port(ctx):
# def batch_results(mgr, results, qname='output'):
# def terminate(mgr, qname='input'):
#   def __init__(self, mgr, train_mode=True, qname_in='input', qname_out='output', input_mapping=None):
#   def next_batch(self, batch_size):
#   def should_stop(self):
#   def batch_results(self, results):
#   def terminate(self):
# class DataFeed(object):


import getpass
import os
import unittest
from tensorflowonspark import TFManager, TFNode
class TFNodeTest(unittest.TestCase):
  def test_hdfs_path(self):
    """Normalization of absolution & relative string paths depending on filesystem"""
    cwd = os.getcwd()
    user = getpass.getuser()
    fs = ["file://", "hdfs://", "viewfs://"]
    paths = {
      "hdfs://foo/bar": ["hdfs://foo/bar", "hdfs://foo/bar", "hdfs://foo/bar"],
      "viewfs://foo/bar": ["viewfs://foo/bar", "viewfs://foo/bar", "viewfs://foo/bar"],
      "file://foo/bar": ["file://foo/bar", "file://foo/bar", "file://foo/bar"],
      "/foo/bar": ["file:///foo/bar", "hdfs:///foo/bar", "viewfs:///foo/bar"],
      "foo/bar": ["file://{}/foo/bar".format(cwd), "hdfs:///user/{}/foo/bar".format(user), "viewfs:///user/{}/foo/bar".format(user)],
    }

    for i in range(len(fs)):
      ctx = type('MockContext', (), {'defaultFS': fs[i], 'working_dir': cwd})
      for path, expected in paths.items():
        final_path = TFNode.hdfs_path(ctx, path)
        self.assertEqual(final_path, expected[i], "fs({}) + path({}) => {}, expected {}".format(fs[i], path, final_path, expected[i]))

  def test_datafeed(self):
    """TFNode.DataFeed basic operations"""
    mgr = TFManager.start('abc'.encode('utf-8'), ['input', 'output'], 'local')

    # insert 10 numbers followed by an end-of-feed marker
    q = mgr.get_queue('input')
    for i in range(10):
      q.put(i)
    q.put(None)                           # end-of-feed marker


# Path: tensorflowonspark/TFManager.py
# class TFManager(BaseManager):
#   """Python multiprocessing.Manager for distributed, multi-process communication."""
#   pass
# 
# Path: tensorflowonspark/TFNode.py
# TF_VERSION = pkg_resources.get_distribution('tensorflow').version
#   HADOOP_SCHEMES = ['adl://',
#                     'file://',
#                     'hdfs://',
#                     'oss://',
#                     's3://',
#                     's3a://',
#                     's3n://',
#                     'swift://',
#                     'viewfs://',
#                     'wasb://']
# def hdfs_path(ctx, path):
# def start_cluster_server(ctx, num_gpus=1, rdma=False):
# def next_batch(mgr, batch_size, qname='input'):
# def export_saved_model(sess, export_dir, tag_set, signatures):
# def release_port(ctx):
# def batch_results(mgr, results, qname='output'):
# def terminate(mgr, qname='input'):
#   def __init__(self, mgr, train_mode=True, qname_in='input', qname_out='output', input_mapping=None):
#   def next_batch(self, batch_size):
#   def should_stop(self):
#   def batch_results(self, results):
#   def terminate(self):
# class DataFeed(object):


# Path: tests/test_TFNode.py
import getpass
import os
import unittest
from tensorflowonspark import TFManager, TFNode

class TFNodeTest(unittest.TestCase):
  def test_hdfs_path(self):
    """Normalization of absolution & relative string paths depending on filesystem"""
    cwd = os.getcwd()
    user = getpass.getuser()
    fs = ["file://", "hdfs://", "viewfs://"]
    paths = {
      "hdfs://foo/bar": ["hdfs://foo/bar", "hdfs://foo/bar", "hdfs://foo/bar"],
      "viewfs://foo/bar": ["viewfs://foo/bar", "viewfs://foo/bar", "viewfs://foo/bar"],
      "file://foo/bar": ["file://foo/bar", "file://foo/bar", "file://foo/bar"],
      "/foo/bar": ["file:///foo/bar", "hdfs:///foo/bar", "viewfs:///foo/bar"],
      "foo/bar": ["file://{}/foo/bar".format(cwd), "hdfs:///user/{}/foo/bar".format(user), "viewfs:///user/{}/foo/bar".format(user)],
    }

    for i in range(len(fs)):
      ctx = type('MockContext', (), {'defaultFS': fs[i], 'working_dir': cwd})
      for path, expected in paths.items():
        final_path = TFNode.hdfs_path(ctx, path)
        self.assertEqual(final_path, expected[i], "fs({}) + path({}) => {}, expected {}".format(fs[i], path, final_path, expected[i]))

  def test_datafeed(self):
    """TFNode.DataFeed basic operations"""
    mgr = TFManager.start('abc'.encode('utf-8'), ['input', 'output'], 'local')

    # insert 10 numbers followed by an end-of-feed marker
    q = mgr.get_queue('input')
    for i in range(10):
      q.put(i)
    q.put(None)                           # end-of-feed marker


    feed = TFNode.DataFeed(mgr)