# Copyright 2018 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.
# =============================================================================

"""Base class for RpcOp tests."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools

import numpy as np

from tensorflow.contrib.proto.python.ops import decode_proto_op
from tensorflow.contrib.proto.python.ops import encode_proto_op
from tensorflow.contrib.rpc.python.kernel_tests import test_example_pb2
from tensorflow.contrib.rpc.python.ops import rpc_op
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors

__all__ = ['I_WARNED_YOU', 'RpcOpTestBase']

I_WARNED_YOU = 'I warned you!'


class RpcOpTestBase(object):
  # pylint: disable=missing-docstring,invalid-name
  """Base class for RpcOp tests."""

  def get_method_name(self, suffix):
    raise NotImplementedError

  def rpc(self, *args, **kwargs):
    return rpc_op.rpc(*args, protocol=self._protocol, **kwargs)

  def try_rpc(self, *args, **kwargs):
    return rpc_op.try_rpc(*args, protocol=self._protocol, **kwargs)

  def testScalarHostPortRpc(self):
    with self.test_session() as sess:
      request_tensors = (
          test_example_pb2.TestCase(shape=[1, 2, 3]).SerializeToString())
      response_tensors = self.rpc(
          method=self.get_method_name('IncrementTestShapes'),
          address=self._address,
          request=request_tensors)
      self.assertEqual(response_tensors.shape, ())
      response_values = sess.run(response_tensors)
    response_message = test_example_pb2.TestCase()
    self.assertTrue(response_message.ParseFromString(response_values))
    self.assertAllEqual([2, 3, 4], response_message.shape)

  def testScalarHostPortTryRpc(self):
    with self.test_session() as sess:
      request_tensors = (
          test_example_pb2.TestCase(shape=[1, 2, 3]).SerializeToString())
      response_tensors, status_code, status_message = self.try_rpc(
          method=self.get_method_name('IncrementTestShapes'),
          address=self._address,
          request=request_tensors)
      self.assertEqual(status_code.shape, ())
      self.assertEqual(status_message.shape, ())
      self.assertEqual(response_tensors.shape, ())
      response_values, status_code_values, status_message_values = (
          sess.run((response_tensors, status_code, status_message)))
    response_message = test_example_pb2.TestCase()
    self.assertTrue(response_message.ParseFromString(response_values))
    self.assertAllEqual([2, 3, 4], response_message.shape)
    # For the base Rpc op, don't expect to get error status back.
    self.assertEqual(errors.OK, status_code_values)
    self.assertEqual(b'', status_message_values)

  def testEmptyHostPortRpc(self):
    with self.test_session() as sess:
      request_tensors = []
      response_tensors = self.rpc(
          method=self.get_method_name('IncrementTestShapes'),
          address=self._address,
          request=request_tensors)
      self.assertAllEqual(response_tensors.shape, [0])
      response_values = sess.run(response_tensors)
    self.assertAllEqual(response_values.shape, [0])

  def testInvalidAddresses(self):
    with self.test_session() as sess:
      with self.assertRaisesOpError(self.invalid_method_string):
        sess.run(
            self.rpc(
                method='/InvalidService.IncrementTestShapes',
                address=self._address,
                request=''))

      with self.assertRaisesOpError(self.invalid_method_string):
        sess.run(
            self.rpc(
                method=self.get_method_name('InvalidMethodName'),
                address=self._address,
                request=''))

      # This also covers the case of address=''
      # and address='localhost:293874293874'
      with self.assertRaises(errors.UnavailableError):
        sess.run(
            self.rpc(
                method=self.get_method_name('IncrementTestShapes'),
                address='unix:/tmp/this_unix_socket_doesnt_exist_97820348!!@',
                request=''))

      # Test invalid method with the TryRpc op
      _, status_code_value, status_message_value = sess.run(
          self.try_rpc(
              method=self.get_method_name('InvalidMethodName'),
              address=self._address,
              request=''))
      self.assertEqual(errors.UNIMPLEMENTED, status_code_value)
      self.assertTrue(
          self.invalid_method_string in status_message_value.decode('ascii'))

  def testAlwaysFailingMethod(self):
    with self.test_session() as sess:
      response_tensors = self.rpc(
          method=self.get_method_name('AlwaysFailWithInvalidArgument'),
          address=self._address,
          request='')
      self.assertEqual(response_tensors.shape, ())
      with self.assertRaisesOpError(I_WARNED_YOU):
        sess.run(response_tensors)

  def testSometimesFailingMethodWithManyRequests(self):
    with self.test_session() as sess:
      # Fail hard by default.
      response_tensors = self.rpc(
          method=self.get_method_name('SometimesFailWithInvalidArgument'),
          address=self._address,
          request=[''] * 20)
      self.assertEqual(response_tensors.shape, (20,))
      with self.assertRaisesOpError(I_WARNED_YOU):
        sess.run(response_tensors)

      # Don't fail hard, use TryRpc - return the failing status instead.
      response_tensors, status_code, status_message = self.try_rpc(
          method=self.get_method_name('SometimesFailWithInvalidArgument'),
          address=self._address,
          request=[''] * 20)
      self.assertEqual(response_tensors.shape, (20,))
      self.assertEqual(status_code.shape, (20,))
      self.assertEqual(status_message.shape, (20,))
      status_code_values, status_message_values = sess.run((status_code,
                                                            status_message))
      self.assertTrue([
          x in (errors.OK, errors.INVALID_ARGUMENT) for x in status_code_values
      ])
      expected_message_values = np.where(
          status_code_values == errors.INVALID_ARGUMENT,
          I_WARNED_YOU.encode('ascii'), b'')
      self.assertAllEqual(expected_message_values, status_message_values)

  def testVecHostPortRpc(self):
    with self.test_session() as sess:
      request_tensors = [
          test_example_pb2.TestCase(
              shape=[i, i + 1, i + 2]).SerializeToString() for i in range(20)
      ]
      response_tensors = self.rpc(
          method=self.get_method_name('IncrementTestShapes'),
          address=self._address,
          request=request_tensors)
      self.assertEqual(response_tensors.shape, (20,))
      response_values = sess.run(response_tensors)
    self.assertEqual(response_values.shape, (20,))
    for i in range(20):
      response_message = test_example_pb2.TestCase()
      self.assertTrue(response_message.ParseFromString(response_values[i]))
      self.assertAllEqual([i + 1, i + 2, i + 3], response_message.shape)

  def testVecHostPortManyParallelRpcs(self):
    with self.test_session() as sess:
      request_tensors = [
          test_example_pb2.TestCase(
              shape=[i, i + 1, i + 2]).SerializeToString() for i in range(20)
      ]
      many_response_tensors = [
          self.rpc(
              method=self.get_method_name('IncrementTestShapes'),
              address=self._address,
              request=request_tensors) for _ in range(10)
      ]
      # Launch parallel 10 calls to the RpcOp, each containing
      # 20 rpc requests.
      many_response_values = sess.run(many_response_tensors)
    self.assertEqual(10, len(many_response_values))
    for response_values in many_response_values:
      self.assertEqual(response_values.shape, (20,))
      for i in range(20):
        response_message = test_example_pb2.TestCase()
        self.assertTrue(response_message.ParseFromString(response_values[i]))
        self.assertAllEqual([i + 1, i + 2, i + 3], response_message.shape)

  def testVecHostPortRpcUsingEncodeAndDecodeProto(self):
    with self.test_session() as sess:
      request_tensors = encode_proto_op.encode_proto(
          message_type='tensorflow.contrib.rpc.TestCase',
          field_names=['shape'],
          sizes=[[3]] * 20,
          values=[
              [[i, i + 1, i + 2] for i in range(20)],
          ])
      response_tensor_strings = self.rpc(
          method=self.get_method_name('IncrementTestShapes'),
          address=self._address,
          request=request_tensors)
      _, (response_shape,) = decode_proto_op.decode_proto(
          bytes=response_tensor_strings,
          message_type='tensorflow.contrib.rpc.TestCase',
          field_names=['shape'],
          output_types=[dtypes.int32])
      response_shape_values = sess.run(response_shape)
    self.assertAllEqual([[i + 1, i + 2, i + 3]
                         for i in range(20)], response_shape_values)

  def testVecHostPortRpcCancelsUponSessionTimeOutWhenSleepingForever(self):
    with self.test_session() as sess:
      request_tensors = [''] * 25  # This will launch 25 RPC requests.
      response_tensors = self.rpc(
          method=self.get_method_name('SleepForever'),
          address=self._address,
          request=request_tensors)
      for timeout_ms in [1, 500, 1000]:
        options = config_pb2.RunOptions(timeout_in_ms=timeout_ms)
        with self.assertRaises((errors.UnavailableError,
                                errors.DeadlineExceededError)):
          sess.run(response_tensors, options=options)

  def testVecHostPortRpcCancelsUponConfiguredTimeOutWhenSleepingForever(self):
    with self.test_session() as sess:
      request_tensors = [''] * 25  # This will launch 25 RPC requests.
      response_tensors = self.rpc(
          method=self.get_method_name('SleepForever'),
          address=self._address,
          timeout_in_ms=1000,
          request=request_tensors)
      with self.assertRaises(errors.DeadlineExceededError):
        sess.run(response_tensors)

  def testTryRpcPropagatesDeadlineErrorWithSometimesTimingOutRequests(self):
    with self.test_session() as sess:
      response_tensors, status_code, status_message = self.try_rpc(
          method=self.get_method_name('SometimesSleepForever'),
          timeout_in_ms=1000,
          address=self._address,
          request=[''] * 20)
      self.assertEqual(response_tensors.shape, (20,))
      self.assertEqual(status_code.shape, (20,))
      self.assertEqual(status_message.shape, (20,))
      status_code_values = sess.run(status_code)
      self.assertTrue([
          x in (errors.OK, errors.DEADLINE_EXCEEDED) for x in status_code_values
      ])

  def testTryRpcWithMultipleAddressesSingleRequest(self):
    flatten = lambda x: list(itertools.chain.from_iterable(x))
    with self.test_session() as sess:
      addresses = flatten([[
          self._address, 'unix:/tmp/this_unix_socket_doesnt_exist_97820348!!@'
      ] for _ in range(10)])
      request = test_example_pb2.TestCase(shape=[0, 1, 2]).SerializeToString()
      response_tensors, status_code, _ = self.try_rpc(
          method=self.get_method_name('IncrementTestShapes'),
          address=addresses,
          request=request)
      response_tensors_values, status_code_values = sess.run((response_tensors,
                                                              status_code))
      self.assertAllEqual(
          flatten([errors.OK, errors.UNAVAILABLE] for _ in range(10)),
          status_code_values)
      for i in range(10):
        self.assertTrue(response_tensors_values[2 * i])
        self.assertFalse(response_tensors_values[2 * i + 1])

  def testTryRpcWithMultipleMethodsSingleRequest(self):
    flatten = lambda x: list(itertools.chain.from_iterable(x))
    with self.test_session() as sess:
      methods = flatten(
          [[self.get_method_name('IncrementTestShapes'), 'InvalidMethodName']
           for _ in range(10)])
      request = test_example_pb2.TestCase(shape=[0, 1, 2]).SerializeToString()
      response_tensors, status_code, _ = self.try_rpc(
          method=methods, address=self._address, request=request)
      response_tensors_values, status_code_values = sess.run((response_tensors,
                                                              status_code))
      self.assertAllEqual(
          flatten([errors.OK, errors.UNIMPLEMENTED] for _ in range(10)),
          status_code_values)
      for i in range(10):
        self.assertTrue(response_tensors_values[2 * i])
        self.assertFalse(response_tensors_values[2 * i + 1])

  def testTryRpcWithMultipleAddressesAndRequests(self):
    flatten = lambda x: list(itertools.chain.from_iterable(x))
    with self.test_session() as sess:
      addresses = flatten([[
          self._address, 'unix:/tmp/this_unix_socket_doesnt_exist_97820348!!@'
      ] for _ in range(10)])
      requests = [
          test_example_pb2.TestCase(
              shape=[i, i + 1, i + 2]).SerializeToString() for i in range(20)
      ]
      response_tensors, status_code, _ = self.try_rpc(
          method=self.get_method_name('IncrementTestShapes'),
          address=addresses,
          request=requests)
      response_tensors_values, status_code_values = sess.run((response_tensors,
                                                              status_code))
      self.assertAllEqual(
          flatten([errors.OK, errors.UNAVAILABLE] for _ in range(10)),
          status_code_values)
      for i in range(20):
        if i % 2 == 1:
          self.assertFalse(response_tensors_values[i])
        else:
          response_message = test_example_pb2.TestCase()
          self.assertTrue(
              response_message.ParseFromString(response_tensors_values[i]))
          self.assertAllEqual([i + 1, i + 2, i + 3], response_message.shape)
