# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for tf.GrpcServer."""

import time

import numpy as np

from xla.tsl.protobuf import rpc_options_pb2
from tensorflow.core.protobuf import cluster_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import tensorflow_server_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_v1
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import input as input_ops
from tensorflow.python.training import queue_runner_impl
from tensorflow.python.training import server_lib


class GrpcServerTest(test.TestCase):

  def __init__(self, methodName="runTest"):  # pylint: disable=invalid-name
    super(GrpcServerTest, self).__init__(methodName)
    self._cached_server = server_lib.Server.create_local_server()

  def testRunStep(self):
    server = self._cached_server
    with ops.Graph().as_default():
      with session.Session(server.target) as sess:
        c = constant_op.constant([[2, 1]])
        d = constant_op.constant([[1], [2]])
        e = math_ops.matmul(c, d)
        self.assertAllEqual([[4]], sess.run(e))
      # TODO(mrry): Add `server.stop()` and `server.join()` when these work.

  def testMultipleSessions(self):
    server = self._cached_server
    with ops.Graph().as_default():
      c = constant_op.constant([[2, 1]])
      d = constant_op.constant([[1], [2]])
      e = math_ops.matmul(c, d)

      sess_1 = session.Session(server.target)
      sess_2 = session.Session(server.target)

      self.assertAllEqual([[4]], sess_1.run(e))
      self.assertAllEqual([[4]], sess_2.run(e))

      sess_1.close()
      sess_2.close()
      # TODO(mrry): Add `server.stop()` and `server.join()` when these work.

  # Verifies various reset failures.
  def testResetFails(self):
    with ops.Graph().as_default():
      # Creates variable with container name.
      with ops.container("test0"):
        v0 = variable_v1.VariableV1(1.0, name="v0")
      # Creates variable with default container.
      v1 = variable_v1.VariableV1(2.0, name="v1")
      # Verifies resetting the non-existent target returns error.
      with self.assertRaises(errors_impl.NotFoundError):
        session.Session.reset("nonexistent", ["test0"])

      # Verifies resetting with config.
      # Verifies that resetting target with no server times out.
      with self.assertRaises(errors_impl.DeadlineExceededError):
        session.Session.reset(
            "grpc://localhost:0",
            ["test0"],
            config=config_pb2.ConfigProto(operation_timeout_in_ms=5),
        )

      # Verifies no containers are reset with non-existent container.
      server = self._cached_server
      sess = session.Session(server.target)
      sess.run(variables.global_variables_initializer())
      self.assertAllEqual(1.0, sess.run(v0))
      self.assertAllEqual(2.0, sess.run(v1))
      # No container is reset, but the server is reset.
      session.Session.reset(server.target, ["test1"])
      # Verifies that both variables are still valid.
      sess = session.Session(server.target)
      self.assertAllEqual(1.0, sess.run(v0))
      self.assertAllEqual(2.0, sess.run(v1))

  def _useRPCConfig(self):
    """Return a `tf.compat.v1.ConfigProto` that ensures we use the RPC stack for tests.

    This configuration ensures that we continue to exercise the gRPC
    stack when testing, rather than using the in-process optimization,
    which avoids using gRPC as the transport between a client and
    master in the same process.

    Returns:
      A `tf.compat.v1.ConfigProto`.
    """
    return config_pb2.ConfigProto(
        rpc_options=rpc_options_pb2.RPCOptions(
            use_rpc_for_inprocess_master=True
        )
    )

  def testLargeConstant(self):
    server = self._cached_server
    with session.Session(server.target, config=self._useRPCConfig()) as sess:
      const_val = np.empty([10000, 3000], dtype=np.float32)
      const_val.fill(0.5)
      c = constant_op.constant(const_val)
      shape_t = array_ops.shape(c)
      self.assertAllEqual([10000, 3000], sess.run(shape_t))

  def testLargeFetch(self):
    server = self._cached_server
    with session.Session(server.target, config=self._useRPCConfig()) as sess:
      c = array_ops.fill([10000, 3000], 0.5)
      expected_val = np.empty([10000, 3000], dtype=np.float32)
      expected_val.fill(0.5)
      self.assertAllEqual(expected_val, sess.run(c))

  def testLargeFeed(self):
    server = self._cached_server
    with session.Session(server.target, config=self._useRPCConfig()) as sess:
      feed_val = np.empty([10000, 3000], dtype=np.float32)
      feed_val.fill(0.5)
      p = array_ops.placeholder(dtypes.float32, shape=[10000, 3000])
      min_t = math_ops.reduce_min(p)
      max_t = math_ops.reduce_max(p)
      min_val, max_val = sess.run([min_t, max_t], feed_dict={p: feed_val})
      self.assertEqual(0.5, min_val)
      self.assertEqual(0.5, max_val)

  def testCloseCancelsBlockingOperation(self):
    server = self._cached_server
    with ops.Graph().as_default():
      sess = session.Session(server.target, config=self._useRPCConfig())

      q = data_flow_ops.FIFOQueue(10, [dtypes.float32])
      enqueue_op = q.enqueue(37.0)
      dequeue_t = q.dequeue()

      sess.run(enqueue_op)
      sess.run(dequeue_t)

      def blocking_dequeue():
        with self.assertRaisesRegex(
            errors_impl.CancelledError, "Session::Close"
        ):
          sess.run(dequeue_t)

      blocking_thread = self.checkedThread(blocking_dequeue)
      blocking_thread.start()
      time.sleep(0.5)
      sess.close()
      blocking_thread.join()

  def testInteractiveSession(self):
    server = self._cached_server
    # Session creation will warn (in C++) that the place_pruned_graph option
    # is not supported, but it should successfully ignore it.
    sess = session.InteractiveSession(server.target)
    c = constant_op.constant(42.0)
    self.assertEqual(42.0, self.evaluate(c))
    sess.close()

  def testSetConfiguration(self):
    config = config_pb2.ConfigProto(
        gpu_options=config_pb2.GPUOptions(per_process_gpu_memory_fraction=0.1)
    )

    # Configure a server using the default local server options.
    server = server_lib.Server.create_local_server(config=config, start=False)
    self.assertEqual(
        0.1,
        server.server_def.default_session_config.gpu_options.per_process_gpu_memory_fraction,
    )

    # Configure a server using an explicit ServerDefd with an
    # overridden config.
    cluster_def = server_lib.ClusterSpec(
        {"localhost": ["localhost:0"]}
    ).as_cluster_def()
    server_def = tensorflow_server_pb2.ServerDef(
        cluster=cluster_def, job_name="localhost", task_index=0, protocol="grpc"
    )
    server = server_lib.Server(server_def, config=config, start=False)
    self.assertEqual(
        0.1,
        server.server_def.default_session_config.gpu_options.per_process_gpu_memory_fraction,
    )

  def testRestartedMaster(self):
    master_old = server_lib.Server.create_local_server()
    master_new = server_lib.Server.create_local_server()
    worker = self._cached_server

    def get_cluster_def(master, worker):
      cluster_def = cluster_pb2.ClusterDef()
      job = cluster_def.job.add()
      job.name = "master"
      job.tasks[0] = master.target[len("grpc://") :]
      job = cluster_def.job.add()
      job.name = "worker"
      job.tasks[0] = worker.target[len("grpc://") :]
      return cluster_def

    def check_session_devices(sess):
      # Make sure we have the correct set of cluster devices
      devices = sess.list_devices()
      device_names = set(d.name for d in devices)
      self.assertIn("/job:master/replica:0/task:0/device:CPU:0", device_names)
      self.assertIn("/job:worker/replica:0/task:0/device:CPU:0", device_names)

    with ops.Graph().as_default():
      # Construct a simple graph that runs ops on remote worker
      with ops.device("/job:worker/replica:0/task:0/device:CPU:0"):
        a = constant_op.constant([1.0])
        b = a + a

      config = config_pb2.ConfigProto(
          cluster_def=get_cluster_def(master_old, worker)
      )
      sess_old = session.Session(master_old.target, config=config)
      check_session_devices(sess_old)

      # Create a session with the new master and the worker.
      # The new master has the same task name ('/job:master/replica:0/task:0')
      # as the old master, but is initiated from a different server thus has a
      # different incarnation. This triggers the WorkerSession on worker with
      # the old master incarnation to be garbage collected.

      config = config_pb2.ConfigProto(
          cluster_def=get_cluster_def(master_new, worker)
      )
      sess_new = session.Session(master_new.target, config=config)
      check_session_devices(sess_new)

      # Running on worker with the new session should work as expected
      v = sess_new.run(b)
      self.assertAllEqual(v, [2.0])

      # Running on worker with the old session should raise an exception since
      # the WorkerSession of the old session has been garbage collected
      with self.assertRaisesRegex(
          errors_impl.AbortedError, "Session handle is not found"
      ):
        sess_old.run(b)

    sess_old.close()
    sess_new.close()

  def testInvalidHostname(self):
    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "port"):
      _ = server_lib.Server(
          {"local": ["localhost"]}, job_name="local", task_index=0
      )

  def testTimeoutRaisesException(self):
    server = self._cached_server
    with ops.Graph().as_default():
      q = data_flow_ops.FIFOQueue(1, [dtypes.float32])
      blocking_t = q.dequeue()

      with session.Session(server.target) as sess:
        with self.assertRaises(errors_impl.DeadlineExceededError):
          sess.run(
              blocking_t, options=config_pb2.RunOptions(timeout_in_ms=1000)
          )

      with session.Session(server.target, config=self._useRPCConfig()) as sess:
        with self.assertRaises(errors_impl.DeadlineExceededError):
          sess.run(
              blocking_t, options=config_pb2.RunOptions(timeout_in_ms=1000)
          )

  def testTwoServersSamePort(self):
    # Starting a server with the same target as the cached server should fail.
    server = self._cached_server
    with self.assertRaises(errors_impl.UnknownError):
      _ = server_lib.Server({"local_2": [server.target[len("grpc://") :]]})

  def testExtendAfterQueueRunners(self):
    server = self._cached_server
    with session.Session(server.target) as sess:
      input_queue = input_ops.input_producer(
          constant_op.constant([0.0], dtype=dtypes.float32)
      )
      self.assertIsNotNone(input_queue)

      var = variable_v1.VariableV1(
          1.0, dtype=dtypes.float32, trainable=False, name="var"
      )

      sess.run(variables.global_variables_initializer())
      queue_runner_impl.start_queue_runners(sess)
      sess.run(var.assign(3.0))

  def testIsolateSessionState(self):
    server = self._cached_server

    with ops.Graph().as_default():
      init_value = array_ops.placeholder(dtypes.int32)
      v = variable_v1.VariableV1(init_value, validate_shape=False, name="v")

      sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
      sharing_sess_0 = session.Session(server.target, config=sharing_config)
      sharing_sess_1 = session.Session(server.target, config=sharing_config)

      isolate_config = config_pb2.ConfigProto(isolate_session_state=True)
      isolate_sess_0 = session.Session(server.target, config=isolate_config)
      isolate_sess_1 = session.Session(server.target, config=isolate_config)

      # Initially all variables are initialized.
      for sess in [
          sharing_sess_0,
          sharing_sess_1,
          isolate_sess_0,
          isolate_sess_1,
      ]:
        with self.assertRaises(errors_impl.FailedPreconditionError):
          sess.run(v)

      # Shared sessions will see each other's updates, but isolated sessions
      # will not.
      sharing_sess_0.run(v.initializer, feed_dict={init_value: 86})
      self.assertAllEqual(86, sharing_sess_0.run(v))
      self.assertAllEqual(86, sharing_sess_1.run(v))
      with self.assertRaises(errors_impl.FailedPreconditionError):
        isolate_sess_0.run(v)
      with self.assertRaises(errors_impl.FailedPreconditionError):
        isolate_sess_1.run(v)

      # Changing the shape works because `validate_shape` is False.
      sharing_sess_1.run(v.initializer, feed_dict={init_value: [86, 99]})
      self.assertAllEqual([86, 99], sharing_sess_0.run(v))
      self.assertAllEqual([86, 99], sharing_sess_1.run(v))
      with self.assertRaises(errors_impl.FailedPreconditionError):
        isolate_sess_0.run(v)
      with self.assertRaises(errors_impl.FailedPreconditionError):
        isolate_sess_1.run(v)

      # Initializing in an isolated session will only affect the state in that
      # session.
      isolate_sess_0.run(v.initializer, feed_dict={init_value: 37})
      self.assertAllEqual([86, 99], sharing_sess_0.run(v))
      self.assertAllEqual([86, 99], sharing_sess_1.run(v))
      self.assertAllEqual(37, isolate_sess_0.run(v))
      with self.assertRaises(errors_impl.FailedPreconditionError):
        isolate_sess_1.run(v)

      # Isolated sessions can have different shapes for the same variable.
      isolate_sess_1.run(v.initializer, feed_dict={init_value: [19, 86]})
      self.assertAllEqual([86, 99], sharing_sess_0.run(v))
      self.assertAllEqual([86, 99], sharing_sess_1.run(v))
      self.assertAllEqual(37, isolate_sess_0.run(v))
      self.assertAllEqual([19, 86], isolate_sess_1.run(v))

  def testShapeChangingIsolateState(self):
    server = self._cached_server
    sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
    isolate_config = config_pb2.ConfigProto(isolate_session_state=True)

    with ops.Graph().as_default():
      w_vector = variable_v1.VariableV1([1, 2, 3], name="w")
      with session.Session(server.target, config=sharing_config) as sess:
        with self.assertRaises(errors_impl.FailedPreconditionError):
          sess.run(w_vector)
        sess.run(w_vector.initializer)
        self.assertAllEqual([1, 2, 3], sess.run(w_vector))

    with ops.Graph().as_default():
      w_vector = variable_v1.VariableV1([4, 5, 6], name="w")
      with session.Session(server.target, config=sharing_config) as sess:
        self.assertAllEqual([1, 2, 3], sess.run(w_vector))
        sess.run(w_vector.initializer)
        self.assertAllEqual([4, 5, 6], sess.run(w_vector))

    with ops.Graph().as_default():
      w_scalar = variable_v1.VariableV1(37, name="w")
      with session.Session(server.target, config=isolate_config) as sess:
        with self.assertRaises(errors_impl.FailedPreconditionError):
          sess.run(w_scalar)
        sess.run(w_scalar.initializer)
        self.assertAllEqual(37, sess.run(w_scalar))


class ServerDefTest(test.TestCase):

  def testLocalServer(self):
    cluster_def = server_lib.ClusterSpec(
        {"local": ["localhost:2222"]}
    ).as_cluster_def()
    server_def = tensorflow_server_pb2.ServerDef(
        cluster=cluster_def, job_name="local", task_index=0, protocol="grpc"
    )

    self.assertProtoEquals(
        """
    cluster {
      job { name: 'local' tasks { key: 0 value: 'localhost:2222' } }
    }
    job_name: 'local' task_index: 0 protocol: 'grpc'
    """,
        server_def,
    )

    # Verifies round trip from Proto->Spec->Proto is correct.
    cluster_spec = server_lib.ClusterSpec(cluster_def)
    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())

  def testTwoProcesses(self):
    cluster_def = server_lib.ClusterSpec(
        {"local": ["localhost:2222", "localhost:2223"]}
    ).as_cluster_def()
    server_def = tensorflow_server_pb2.ServerDef(
        cluster=cluster_def, job_name="local", task_index=1, protocol="grpc"
    )

    self.assertProtoEquals(
        """
    cluster {
      job { name: 'local' tasks { key: 0 value: 'localhost:2222' }
                          tasks { key: 1 value: 'localhost:2223' } }
    }
    job_name: 'local' task_index: 1 protocol: 'grpc'
    """,
        server_def,
    )

    # Verifies round trip from Proto->Spec->Proto is correct.
    cluster_spec = server_lib.ClusterSpec(cluster_def)
    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())

  def testTwoJobs(self):
    cluster_def = server_lib.ClusterSpec({
        "ps": ["ps0:2222", "ps1:2222"],
        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
    }).as_cluster_def()
    server_def = tensorflow_server_pb2.ServerDef(
        cluster=cluster_def, job_name="worker", task_index=2, protocol="grpc"
    )

    self.assertProtoEquals(
        """
    cluster {
      job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
                       tasks { key: 1 value: 'ps1:2222' } }
      job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
                           tasks { key: 1 value: 'worker1:2222' }
                           tasks { key: 2 value: 'worker2:2222' } }
    }
    job_name: 'worker' task_index: 2 protocol: 'grpc'
    """,
        server_def,
    )

    # Verifies round trip from Proto->Spec->Proto is correct.
    cluster_spec = server_lib.ClusterSpec(cluster_def)
    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())

  def testDenseAndSparseJobs(self):
    cluster_def = server_lib.ClusterSpec({
        "ps": ["ps0:2222", "ps1:2222"],
        "worker": {0: "worker0:2222", 2: "worker2:2222"},
    }).as_cluster_def()
    server_def = tensorflow_server_pb2.ServerDef(
        cluster=cluster_def, job_name="worker", task_index=2, protocol="grpc"
    )

    self.assertProtoEquals(
        """
    cluster {
      job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
                       tasks { key: 1 value: 'ps1:2222' } }
      job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
                           tasks { key: 2 value: 'worker2:2222' } }
    }
    job_name: 'worker' task_index: 2 protocol: 'grpc'
    """,
        server_def,
    )

    # Verifies round trip from Proto->Spec->Proto is correct.
    cluster_spec = server_lib.ClusterSpec(cluster_def)
    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())


class ClusterSpecTest(test.TestCase):

  def testStringConversion(self):
    cluster_spec = server_lib.ClusterSpec(
        {"ps": ["ps0:1111"], "worker": ["worker0:3333", "worker1:4444"]}
    )

    expected_str = (
        "ClusterSpec({'ps': ['ps0:1111'], 'worker': ['worker0:3333', "
        "'worker1:4444']})"
    )
    self.assertEqual(expected_str, str(cluster_spec))

  def testProtoDictDefEquivalences(self):
    cluster_spec = server_lib.ClusterSpec({
        "ps": ["ps0:2222", "ps1:2222"],
        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
    })

    expected_proto = """
    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
                     tasks { key: 1 value: 'ps1:2222' } }
    job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
                         tasks { key: 1 value: 'worker1:2222' }
                         tasks { key: 2 value: 'worker2:2222' } }
    """

    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
    self.assertProtoEquals(
        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
    )
    self.assertProtoEquals(
        expected_proto,
        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
    )
    self.assertProtoEquals(
        expected_proto,
        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
    )

  def testProtoDictDefEquivalencesWithStringTaskIndex(self):
    cluster_spec = server_lib.ClusterSpec(
        {"ps": ["ps0:2222", "ps1:2222"], "worker": {"1": "worker1:2222"}}
    )

    expected_proto = """
    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
                     tasks { key: 1 value: 'ps1:2222' } }
    job { name: 'worker' tasks { key: 1 value: 'worker1:2222' } }
    """

    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
    self.assertProtoEquals(
        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
    )
    self.assertProtoEquals(
        expected_proto,
        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
    )
    self.assertProtoEquals(
        expected_proto,
        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
    )

  def testProtoDictDefEquivalencesWithZeroWorker(self):
    cluster_spec = server_lib.ClusterSpec(
        {"ps": ["ps0:2222", "ps1:2222"], "worker": []}
    )

    expected_proto = """
    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
                     tasks { key: 1 value: 'ps1:2222' } }
    job { name: 'worker' }
    """

    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
    self.assertProtoEquals(
        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
    )
    self.assertProtoEquals(
        expected_proto,
        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
    )
    self.assertProtoEquals(
        expected_proto,
        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
    )

  def testClusterSpecAccessors(self):
    original_dict = {
        "ps": ["ps0:2222", "ps1:2222"],
        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
        "sparse": {0: "sparse0:2222", 3: "sparse3:2222"},
    }
    cluster_spec = server_lib.ClusterSpec(original_dict)

    self.assertEqual(original_dict, cluster_spec.as_dict())

    self.assertEqual(2, cluster_spec.num_tasks("ps"))
    self.assertEqual(3, cluster_spec.num_tasks("worker"))
    self.assertEqual(2, cluster_spec.num_tasks("sparse"))
    with self.assertRaises(ValueError):
      cluster_spec.num_tasks("unknown")

    self.assertEqual("ps0:2222", cluster_spec.task_address("ps", 0))
    self.assertEqual("sparse0:2222", cluster_spec.task_address("sparse", 0))
    with self.assertRaises(ValueError):
      cluster_spec.task_address("unknown", 0)
    with self.assertRaises(ValueError):
      cluster_spec.task_address("sparse", 2)

    self.assertEqual([0, 1], cluster_spec.task_indices("ps"))
    self.assertEqual([0, 1, 2], cluster_spec.task_indices("worker"))
    self.assertEqual([0, 3], cluster_spec.task_indices("sparse"))
    with self.assertRaises(ValueError):
      cluster_spec.task_indices("unknown")

    # NOTE(mrry): `ClusterSpec.job_tasks()` is not recommended for use
    # with sparse jobs.
    self.assertEqual(["ps0:2222", "ps1:2222"], cluster_spec.job_tasks("ps"))
    self.assertEqual(
        ["worker0:2222", "worker1:2222", "worker2:2222"],
        cluster_spec.job_tasks("worker"),
    )
    self.assertEqual(
        ["sparse0:2222", None, None, "sparse3:2222"],
        cluster_spec.job_tasks("sparse"),
    )
    with self.assertRaises(ValueError):
      cluster_spec.job_tasks("unknown")

  def testEmptyClusterSpecIsFalse(self):
    self.assertFalse(server_lib.ClusterSpec({}))

  def testNonEmptyClusterSpecIsTrue(self):
    self.assertTrue(server_lib.ClusterSpec({"job": ["host:port"]}))

  def testEq(self):
    self.assertEqual(server_lib.ClusterSpec({}), server_lib.ClusterSpec({}))
    self.assertEqual(
        server_lib.ClusterSpec({"job": ["host:2222"]}),
        server_lib.ClusterSpec({"job": ["host:2222"]}),
    )
    self.assertEqual(
        server_lib.ClusterSpec({"job": {0: "host:2222"}}),
        server_lib.ClusterSpec({"job": ["host:2222"]}),
    )

  def testNe(self):
    self.assertNotEqual(
        server_lib.ClusterSpec({}),
        server_lib.ClusterSpec({"job": ["host:2223"]}),
    )
    self.assertNotEqual(
        server_lib.ClusterSpec({"job1": ["host:2222"]}),
        server_lib.ClusterSpec({"job2": ["host:2222"]}),
    )
    self.assertNotEqual(
        server_lib.ClusterSpec({"job": ["host:2222"]}),
        server_lib.ClusterSpec({"job": ["host:2223"]}),
    )
    self.assertNotEqual(
        server_lib.ClusterSpec({"job": ["host:2222", "host:2223"]}),
        server_lib.ClusterSpec({"job": ["host:2223", "host:2222"]}),
    )


if __name__ == "__main__":
  test.main()