# Copyright 2014 MongoDB, Inc.
#
# 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.

"""Utilities for testing Motor with asyncio."""

import asyncio
import functools
import gc
import inspect
import unittest
from asyncio import ensure_future
from test.assert_logs_backport import AssertLogsMixin
from test.test_environment import CA_PEM, CLIENT_PEM, env
from test.utils import get_async_test_timeout
from unittest import SkipTest

from mockupdb import MockupDB

from motor import motor_asyncio

# mypy: ignore-errors


class _TestMethodWrapper:
    """Wraps a test method to raise an error if it returns a value.

    This is mainly used to detect undecorated generators (if a test
    method yields it must use a decorator to consume the generator),
    but will also detect other kinds of return values (these are not
    necessarily errors, but we alert anyway since there is no good
    reason to return a value from a test).

    Adapted from Tornado's test framework.
    """

    def __init__(self, orig_method):
        self.orig_method = orig_method

    def __call__(self):
        result = self.orig_method()
        if inspect.iscoroutine(result):
            # Cancel the undecorated task to avoid this warning:
            # RuntimeWarning: coroutine 'test_foo' was never awaited
            task = ensure_future(result)
            task.cancel()
            raise TypeError("Generator test methods should be decorated with @asyncio_test")
        elif result is not None:
            raise ValueError("Return value from test method ignored: %r" % result)

    def __getattr__(self, name):
        """Proxy all unknown attributes to the original method.

        This is important for some of the decorators in the `unittest`
        module, such as `unittest.skipIf`.
        """
        return getattr(self.orig_method, name)


class AsyncIOTestCase(AssertLogsMixin, unittest.TestCase):
    longMessage = True  # Used by unittest.TestCase
    ssl = False  # If True, connect with SSL, skip if mongod isn't SSL

    def __init__(self, methodName="runTest"):
        super().__init__(methodName)

        # It's easy to forget the @asyncio_test decorator, but if you do
        # the test will silently be ignored because nothing will consume
        # the generator. Replace the test method with a wrapper that will
        # make sure it's not an undecorated generator.
        # (Adapted from Tornado's AsyncTestCase.)
        setattr(self, methodName, _TestMethodWrapper(getattr(self, methodName)))

    def setUp(self):
        super().setUp()

        # Ensure that the event loop is passed explicitly in Motor.
        asyncio.set_event_loop(None)
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        if self.ssl and not env.mongod_started_with_ssl:
            raise SkipTest("mongod doesn't support SSL, or is down")

        self.cx = self.asyncio_client()
        self.db = self.cx.motor_test
        self.collection = self.db.test_collection
        self.loop.run_until_complete(self.collection.drop())

    # Workaround for https://github.com/pytest-dev/pytest/issues/12263.
    def runTest(self):
        pass

    def get_client_kwargs(self, set_loop=True, **kwargs):
        if set_loop:
            kwargs.setdefault("io_loop", self.loop)
        if env.mongod_started_with_ssl:
            kwargs.setdefault("tlsCAFile", CA_PEM)
            kwargs.setdefault("tlsCertificateKeyFile", CLIENT_PEM)
        kwargs.setdefault("tls", env.mongod_started_with_ssl)
        return kwargs

    def asyncio_client(
        self, uri=None, *args, set_loop=True, **kwargs
    ) -> motor_asyncio.AsyncIOMotorClient:
        """Get an AsyncIOMotorClient.

        Ignores self.ssl, you must pass 'ssl' argument.
        """
        return motor_asyncio.AsyncIOMotorClient(
            kwargs.pop("host", None) or uri or env.uri,
            *args,
            **self.get_client_kwargs(**kwargs, set_loop=set_loop),
        )

    def asyncio_rsc(self, uri=None, *args, **kwargs):
        """Get an open MotorClient for replica set.

        Ignores self.ssl, you must pass 'ssl' argument.
        """
        return motor_asyncio.AsyncIOMotorClient(
            uri or env.rs_uri, *args, **self.get_client_kwargs(**kwargs)
        )

    async def make_test_data(self):
        await self.collection.delete_many({})
        await self.collection.insert_many([{"_id": i} for i in range(200)])

    make_test_data.__test__ = False

    def tearDown(self):
        self.cx.close()
        self.loop.stop()
        self.loop.run_forever()
        self.loop.close()
        asyncio.set_event_loop(asyncio.new_event_loop())
        gc.collect()


class AsyncIOMockServerTestCase(AsyncIOTestCase):
    def server(self, *args, **kwargs):
        server = MockupDB(*args, **kwargs)
        server.run()
        self.addCleanup(server.stop)
        return server

    def client_server(self, *args, **kwargs):
        server = self.server(*args, **kwargs)
        client = motor_asyncio.AsyncIOMotorClient(server.uri, io_loop=self.loop)
        self.addCleanup(client.close)

        return client, server

    def run_thread(self, fn, *args, **kwargs):
        return self.loop.run_in_executor(None, functools.partial(fn, *args, **kwargs))

    def fetch_next(self, cursor):
        async def fetch_next():
            return await cursor.fetch_next

        return ensure_future(fetch_next())


# TODO: Spin off to a PyPI package.
def asyncio_test(func=None, timeout=None):
    """Decorator for coroutine methods of AsyncIOTestCase::

        class MyTestCase(AsyncIOTestCase):
            @asyncio_test
            def test(self):
                # Your test code here....
                pass

    Default timeout is 5 seconds. Override like::

        class MyTestCase(AsyncIOTestCase):
            @asyncio_test(timeout=10)
            def test(self):
                # Your test code here....
                pass

    You can also set the ASYNC_TEST_TIMEOUT environment variable to a number
    of seconds. The final timeout is the ASYNC_TEST_TIMEOUT or the timeout
    in the test (5 seconds or the passed-in timeout), whichever is longest.
    """

    def wrap(f):
        @functools.wraps(f)
        def wrapped(self, *args, **kwargs):
            if timeout is None:
                actual_timeout = get_async_test_timeout()
            else:
                actual_timeout = get_async_test_timeout(timeout)

            coro_exc = None

            def exc_handler(loop, context):
                nonlocal coro_exc
                # Exception is optional.
                coro_exc = context.get("exception", Exception(context))

                # Raise CancelledError from run_until_complete below.
                task.cancel()

            self.loop.set_exception_handler(exc_handler)
            coro = asyncio.wait_for(f(self, *args, **kwargs), actual_timeout)
            task = ensure_future(coro, loop=self.loop)
            try:
                self.loop.run_until_complete(task)
            except BaseException:
                if coro_exc:
                    # Raise the error thrown in on_timeout, with only the
                    # traceback from the coroutine itself, not from
                    # run_until_complete.
                    raise coro_exc from None

                raise

        return wrapped

    if func is not None:
        # Used like:
        #     @gen_test
        #     def f(self):
        #         pass
        if not inspect.isfunction(func):
            msg = (
                "%r is not a test method. Pass a timeout as"
                " a keyword argument, like @asyncio_test(timeout=7)"
            )
            raise TypeError(msg % func)
        return wrap(func)
    else:
        # Used like @gen_test(timeout=10)
        return wrap


async def get_command_line(client):
    command_line = await client.admin.command("getCmdLineOpts")
    assert command_line["ok"] == 1, "getCmdLineOpts() failed"
    return command_line


async def server_is_mongos(client):
    ismaster_response = await client.admin.command("ismaster")
    return ismaster_response.get("msg") == "isdbgrid"


async def skip_if_mongos(client):
    is_mongos = await server_is_mongos(client)
    if is_mongos:
        raise unittest.SkipTest("connected to mongos")


async def remove_all_users(db):
    await db.command({"dropAllUsersFromDatabase": 1})