# Path: tensorflowonspark/TFParallel.py
# def run(sc, map_fn, tf_args, num_executors, use_barrier=True):
#   def _run(it):


import unittest
import test
      import tensorflow as tf
      import tensorflow as tf
      import tensorflow as tf
from tensorflowonspark import TFParallel
    def _map_fun(args, ctx):
      x = tf.constant(args['x'])
      y = tf.constant(args['y'])
      sum = tf.math.add(x, y)
      assert sum.numpy() == 3

    args = {'x': 1, 'y': 2}
    TFParallel.run(self.sc, _map_fun, tf_args=args, num_executors=self.num_workers, use_barrier=False)

  def test_basic_tf_barrier(self):
    """Single-node TF graph (w/ args) running independently on multiple executors using Spark barrier."""
    def _map_fun(args, ctx):
      x = tf.constant(args['x'])
      y = tf.constant(args['y'])
      sum = tf.math.add(x, y)
      assert sum.numpy() == 3

    args = {'x': 1, 'y': 2}
    TFParallel.run(self.sc, _map_fun, tf_args=args, num_executors=self.num_workers)

  def test_basic_tf_barrier_insufficient_resources(self):
    """Single-node TF graph (w/ args) running independently on multiple executors using Spark barrier with insufficient resource."""
    def _map_fun(args, ctx):
      x = tf.constant(args['x'])
      y = tf.constant(args['y'])
      sum = tf.math.add(x, y)
      assert sum.numpy() == 3

    args = {'x': 1, 'y': 2}
    with self.assertRaises(Exception):

# Path: tensorflowonspark/TFParallel.py
# def run(sc, map_fn, tf_args, num_executors, use_barrier=True):
#   def _run(it):


# Path: tests/test_TFParallel.py
import unittest
import test
      import tensorflow as tf
      import tensorflow as tf
      import tensorflow as tf
from tensorflowonspark import TFParallel

    def _map_fun(args, ctx):
      x = tf.constant(args['x'])
      y = tf.constant(args['y'])
      sum = tf.math.add(x, y)
      assert sum.numpy() == 3

    args = {'x': 1, 'y': 2}
    TFParallel.run(self.sc, _map_fun, tf_args=args, num_executors=self.num_workers, use_barrier=False)

  def test_basic_tf_barrier(self):
    """Single-node TF graph (w/ args) running independently on multiple executors using Spark barrier."""
    def _map_fun(args, ctx):
      x = tf.constant(args['x'])
      y = tf.constant(args['y'])
      sum = tf.math.add(x, y)
      assert sum.numpy() == 3

    args = {'x': 1, 'y': 2}
    TFParallel.run(self.sc, _map_fun, tf_args=args, num_executors=self.num_workers)

  def test_basic_tf_barrier_insufficient_resources(self):
    """Single-node TF graph (w/ args) running independently on multiple executors using Spark barrier with insufficient resource."""
    def _map_fun(args, ctx):
      x = tf.constant(args['x'])
      y = tf.constant(args['y'])
      sum = tf.math.add(x, y)
      assert sum.numpy() == 3

    args = {'x': 1, 'y': 2}
    with self.assertRaises(Exception):

      TFParallel.run(self.sc, _map_fun, tf_args=args, num_executors=self.num_workers + 1)