# Copyright 2017 gRPC authors.
#
# 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.

from concurrent import futures
import time
import unittest

import grpc
from grpc.framework.foundation import logging_pool
import grpc_testing

from tests.testing import _application_common
from tests.testing import _application_testing_common
from tests.testing import _client_application
from tests.testing.proto import requests_pb2
from tests.testing.proto import services_pb2
from tests.unit.framework.common import test_constants


# TODO(https://github.com/protocolbuffers/protobuf/issues/3452): Drop this skip.
@unittest.skipIf(
    services_pb2.DESCRIPTOR.services_by_name.get('FirstService') is None,
    'Fix protobuf issue 3452!')
class ClientTest(unittest.TestCase):

    def setUp(self):
        # In this test the client-side application under test executes in
        # a separate thread while we retain use of the test thread to "play
        # server".
        self._client_execution_thread_pool = logging_pool.pool(1)

        self._fake_time = grpc_testing.strict_fake_time(time.time())
        self._real_time = grpc_testing.strict_real_time()
        self._fake_time_channel = grpc_testing.channel(
            services_pb2.DESCRIPTOR.services_by_name.values(), self._fake_time)
        self._real_time_channel = grpc_testing.channel(
            services_pb2.DESCRIPTOR.services_by_name.values(), self._real_time)

    def tearDown(self):
        self._client_execution_thread_pool.shutdown(wait=True)

    def test_successful_unary_unary(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run, _client_application.Scenario.UNARY_UNARY,
            self._real_time_channel)
        invocation_metadata, request, rpc = (
            self._real_time_channel.take_unary_unary(
                _application_testing_common.FIRST_SERVICE_UNUN))
        rpc.send_initial_metadata(())
        rpc.terminate(_application_common.UNARY_UNARY_RESPONSE, (),
                      grpc.StatusCode.OK, '')
        application_return_value = application_future.result()

        self.assertEqual(_application_common.UNARY_UNARY_REQUEST, request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.SATISFACTORY)

    def test_successful_unary_stream(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run, _client_application.Scenario.UNARY_STREAM,
            self._fake_time_channel)
        invocation_metadata, request, rpc = (
            self._fake_time_channel.take_unary_stream(
                _application_testing_common.FIRST_SERVICE_UNSTRE))
        rpc.send_initial_metadata(())
        rpc.terminate((), grpc.StatusCode.OK, '')
        application_return_value = application_future.result()

        self.assertEqual(_application_common.UNARY_STREAM_REQUEST, request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.SATISFACTORY)

    def test_successful_stream_unary(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run, _client_application.Scenario.STREAM_UNARY,
            self._real_time_channel)
        invocation_metadata, rpc = self._real_time_channel.take_stream_unary(
            _application_testing_common.FIRST_SERVICE_STREUN)
        rpc.send_initial_metadata(())
        first_request = rpc.take_request()
        second_request = rpc.take_request()
        third_request = rpc.take_request()
        rpc.requests_closed()
        rpc.terminate(_application_common.STREAM_UNARY_RESPONSE, (),
                      grpc.StatusCode.OK, '')
        application_return_value = application_future.result()

        self.assertEqual(_application_common.STREAM_UNARY_REQUEST,
                         first_request)
        self.assertEqual(_application_common.STREAM_UNARY_REQUEST,
                         second_request)
        self.assertEqual(_application_common.STREAM_UNARY_REQUEST,
                         third_request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.SATISFACTORY)

    def test_successful_stream_stream(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run, _client_application.Scenario.STREAM_STREAM,
            self._fake_time_channel)
        invocation_metadata, rpc = self._fake_time_channel.take_stream_stream(
            _application_testing_common.FIRST_SERVICE_STRESTRE)
        first_request = rpc.take_request()
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        second_request = rpc.take_request()
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.requests_closed()
        rpc.terminate((), grpc.StatusCode.OK, '')
        application_return_value = application_future.result()

        self.assertEqual(_application_common.STREAM_STREAM_REQUEST,
                         first_request)
        self.assertEqual(_application_common.STREAM_STREAM_REQUEST,
                         second_request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.SATISFACTORY)

    def test_concurrent_stream_stream(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run,
            _client_application.Scenario.CONCURRENT_STREAM_STREAM,
            self._real_time_channel)
        rpcs = []
        for _ in range(test_constants.RPC_CONCURRENCY):
            invocation_metadata, rpc = (
                self._real_time_channel.take_stream_stream(
                    _application_testing_common.FIRST_SERVICE_STRESTRE))
            rpcs.append(rpc)
        requests = {}
        for rpc in rpcs:
            requests[rpc] = [rpc.take_request()]
        for rpc in rpcs:
            rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
            rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        for rpc in rpcs:
            requests[rpc].append(rpc.take_request())
        for rpc in rpcs:
            rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
            rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        for rpc in rpcs:
            rpc.requests_closed()
        for rpc in rpcs:
            rpc.terminate((), grpc.StatusCode.OK, '')
        application_return_value = application_future.result()

        for requests_of_one_rpc in requests.values():
            for request in requests_of_one_rpc:
                self.assertEqual(_application_common.STREAM_STREAM_REQUEST,
                                 request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.SATISFACTORY)

    def test_cancelled_unary_unary(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run,
            _client_application.Scenario.CANCEL_UNARY_UNARY,
            self._fake_time_channel)
        invocation_metadata, request, rpc = (
            self._fake_time_channel.take_unary_unary(
                _application_testing_common.FIRST_SERVICE_UNUN))
        rpc.send_initial_metadata(())
        rpc.cancelled()
        application_return_value = application_future.result()

        self.assertEqual(_application_common.UNARY_UNARY_REQUEST, request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.SATISFACTORY)

    def test_status_stream_unary(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run,
            _client_application.Scenario.CONCURRENT_STREAM_UNARY,
            self._fake_time_channel)
        rpcs = tuple(
            self._fake_time_channel.take_stream_unary(
                _application_testing_common.FIRST_SERVICE_STREUN)[1]
            for _ in range(test_constants.THREAD_CONCURRENCY))
        for rpc in rpcs:
            rpc.take_request()
            rpc.take_request()
            rpc.take_request()
            rpc.requests_closed()
            rpc.send_initial_metadata(((
                'my_metadata_key',
                'My Metadata Value!',
            ),))
        for rpc in rpcs[:-1]:
            rpc.terminate(_application_common.STREAM_UNARY_RESPONSE, (),
                          grpc.StatusCode.OK, '')
        rpcs[-1].terminate(_application_common.STREAM_UNARY_RESPONSE, (),
                           grpc.StatusCode.RESOURCE_EXHAUSTED,
                           'nope; not able to handle all those RPCs!')
        application_return_value = application_future.result()

        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.UNSATISFACTORY)

    def test_status_stream_stream(self):
        code = grpc.StatusCode.DEADLINE_EXCEEDED
        details = 'test deadline exceeded!'

        application_future = self._client_execution_thread_pool.submit(
            _client_application.run, _client_application.Scenario.STREAM_STREAM,
            self._real_time_channel)
        invocation_metadata, rpc = self._real_time_channel.take_stream_stream(
            _application_testing_common.FIRST_SERVICE_STRESTRE)
        first_request = rpc.take_request()
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        second_request = rpc.take_request()
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.requests_closed()
        rpc.terminate((), code, details)
        application_return_value = application_future.result()

        self.assertEqual(_application_common.STREAM_STREAM_REQUEST,
                         first_request)
        self.assertEqual(_application_common.STREAM_STREAM_REQUEST,
                         second_request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.RPC_ERROR)
        self.assertIs(application_return_value.code, code)
        self.assertEqual(application_return_value.details, details)

    def test_misbehaving_server_unary_unary(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run, _client_application.Scenario.UNARY_UNARY,
            self._fake_time_channel)
        invocation_metadata, request, rpc = (
            self._fake_time_channel.take_unary_unary(
                _application_testing_common.FIRST_SERVICE_UNUN))
        rpc.send_initial_metadata(())
        rpc.terminate(_application_common.ERRONEOUS_UNARY_UNARY_RESPONSE, (),
                      grpc.StatusCode.OK, '')
        application_return_value = application_future.result()

        self.assertEqual(_application_common.UNARY_UNARY_REQUEST, request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.UNSATISFACTORY)

    def test_misbehaving_server_stream_stream(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run, _client_application.Scenario.STREAM_STREAM,
            self._real_time_channel)
        invocation_metadata, rpc = self._real_time_channel.take_stream_stream(
            _application_testing_common.FIRST_SERVICE_STRESTRE)
        first_request = rpc.take_request()
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        second_request = rpc.take_request()
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.send_response(_application_common.STREAM_STREAM_RESPONSE)
        rpc.requests_closed()
        rpc.terminate((), grpc.StatusCode.OK, '')
        application_return_value = application_future.result()

        self.assertEqual(_application_common.STREAM_STREAM_REQUEST,
                         first_request)
        self.assertEqual(_application_common.STREAM_STREAM_REQUEST,
                         second_request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.UNSATISFACTORY)

    def test_infinite_request_stream_real_time(self):
        application_future = self._client_execution_thread_pool.submit(
            _client_application.run,
            _client_application.Scenario.INFINITE_REQUEST_STREAM,
            self._real_time_channel)
        invocation_metadata, rpc = self._real_time_channel.take_stream_unary(
            _application_testing_common.FIRST_SERVICE_STREUN)
        rpc.send_initial_metadata(())
        first_request = rpc.take_request()
        second_request = rpc.take_request()
        third_request = rpc.take_request()
        self._real_time.sleep_for(
            _application_common.INFINITE_REQUEST_STREAM_TIMEOUT)
        rpc.terminate(_application_common.STREAM_UNARY_RESPONSE, (),
                      grpc.StatusCode.DEADLINE_EXCEEDED, '')
        application_return_value = application_future.result()

        self.assertEqual(_application_common.STREAM_UNARY_REQUEST,
                         first_request)
        self.assertEqual(_application_common.STREAM_UNARY_REQUEST,
                         second_request)
        self.assertEqual(_application_common.STREAM_UNARY_REQUEST,
                         third_request)
        self.assertIs(application_return_value.kind,
                      _client_application.Outcome.Kind.SATISFACTORY)


if __name__ == '__main__':
    unittest.main(verbosity=2)
