# Copyright 2016 Quora, 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.

import asynq
from qcore.asserts import assert_eq, assert_is, assert_not_in, AssertRaises

# ===================================================
# Objects to mock
# ===================================================


@asynq.asynq()
def fn():
    pass


@asynq.asynq()
def async_caller():
    ret = yield fn.asynq()
    asynq.result((ret, fn())); return


def non_async_caller():
    return fn()


class Cls(object):
    @asynq.asynq()
    @classmethod
    def async_classmethod(cls):
        pass

    def non_async_method(self, foo):
        pass

    @asynq.asynq()
    def async_method(self):
        return 'capybaras'

instance = Cls()


@asynq.asynq()
def class_method_async_caller():
    ret = yield Cls.async_classmethod.asynq()
    asynq.result((ret, Cls.async_classmethod())); return


def class_method_non_async_caller():
    return Cls.async_classmethod()


@asynq.asynq()
def method_async_caller():
    obj = Cls()
    ret = yield obj.async_method.asynq()
    asynq.result((ret, obj.async_method())); return


def method_non_async_caller():
    return Cls().async_method()


class Counter(object):
    linked_class = Cls

    def __init__(self):
        self.count = 0

    def add_call(self):
        self.count += 1


class SecondClass(object):
    pass


IMPORTANT_DICTIONARY = {'capybaras': 'guinea pigs', 'hutias': 'coypus'}

# ===================================================
# Helpers for testing that mocks work right.
# ===================================================


class MockChecker(object):
    @classmethod
    def check(cls, mock_fn, mock_classmethod, mock_method):
        cls._check_mock(mock_fn, async_caller, non_async_caller)
        cls._check_mock(mock_classmethod, class_method_async_caller, class_method_non_async_caller)
        cls._check_mock(mock_method, method_async_caller, method_non_async_caller)

    @classmethod
    def _check_mock(cls, mock_fn, async_caller, non_async_caller):
        raise NotImplementedError


class MockCheckerWithAssignment(MockChecker):
    @classmethod
    def _check_mock(cls, mock_fn, async_caller, non_async_caller):
        mock_fn.return_value = 42
        assert_eq(0, mock_fn.call_count)
        assert_eq(42, non_async_caller())
        assert_eq(1, mock_fn.call_count)

        assert_eq((42, 42), async_caller())
        assert_eq(3, mock_fn.call_count)

        mock_fn.side_effect = lambda: 43
        assert_eq(43, non_async_caller())
        assert_eq((43, 43), async_caller())


class MockCheckerWithNew(MockChecker):
    @classmethod
    def _check_mock(cls, mock_fn, async_caller, non_async_caller):
        assert_eq(42, non_async_caller())
        assert_eq((42, 42), async_caller())


# ===================================================
# Actual tests.
# ===================================================


def test_mock_async_context():
    with asynq.mock.patch('asynq.tests.test_mock.fn') as mock_fn, \
            asynq.mock.patch.object(Cls, 'async_classmethod') as mock_classmethod, \
            asynq.mock.patch.object(Cls, 'async_method') as mock_method:
        MockCheckerWithAssignment.check(mock_fn, mock_classmethod, mock_method)

    with asynq.mock.patch('asynq.tests.test_mock.fn', lambda: 42) as mock_fn, \
            asynq.mock.patch.object(Cls, 'async_classmethod', classmethod(lambda _: 42)) as mock_classmethod, \
            asynq.mock.patch.object(Cls, 'async_method', lambda _: 42) as mock_method:
        MockCheckerWithNew.check(mock_fn, mock_classmethod, mock_method)


@asynq.mock.patch('asynq.tests.test_mock.fn')
@asynq.mock.patch.object(Cls, 'async_classmethod')
@asynq.mock.patch('asynq.tests.test_mock.Cls.async_method')
def test_mock_async_decorator(mock_method, mock_classmethod, mock_fn):
    MockCheckerWithAssignment.check(mock_fn, mock_classmethod, mock_method)


@asynq.mock.patch('asynq.tests.test_mock.fn')
@asynq.mock.patch.object(Cls, 'async_classmethod')
@asynq.mock.patch('asynq.tests.test_mock.Cls.async_method')
class TestMockAsyncClassDecorator(object):
    def test(self, mock_method, mock_classmethod, mock_fn):
        MockCheckerWithAssignment.check(mock_fn, mock_classmethod, mock_method)


@asynq.mock.patch('asynq.tests.test_mock.fn', lambda: 42)
@asynq.mock.patch.object(Cls, 'async_classmethod', classmethod(lambda _: 42))
@asynq.mock.patch('asynq.tests.test_mock.Cls.async_method', lambda _: 42)
def test_mock_async_decorator_with_new():
    MockCheckerWithNew.check(fn, Cls.async_classmethod, Cls().async_method)


@asynq.mock.patch('asynq.tests.test_mock.fn', lambda: 42)
@asynq.mock.patch.object(Cls, 'async_classmethod', classmethod(lambda _: 42))
@asynq.mock.patch('asynq.tests.test_mock.Cls.async_method', lambda _: 42)
class TestMockAsyncClassDecoratorWithNew(object):
    def test(self):
        MockCheckerWithNew.check(fn, Cls.async_classmethod, Cls().async_method)


def test_mock_async_dict():
    assert_eq('guinea pigs', IMPORTANT_DICTIONARY['capybaras'])

    with asynq.mock.patch.dict(
            'asynq.tests.test_mock.IMPORTANT_DICTIONARY', {'capybaras': 'maras'}):
        assert_eq('maras', IMPORTANT_DICTIONARY['capybaras'])
        assert_eq('coypus', IMPORTANT_DICTIONARY['hutias'])
    assert_eq('guinea pigs', IMPORTANT_DICTIONARY['capybaras'])

    with asynq.mock.patch.dict('asynq.tests.test_mock.IMPORTANT_DICTIONARY', {'capybaras': 'maras'},
                               clear=True):
        assert_eq('maras', IMPORTANT_DICTIONARY['capybaras'])
        assert_not_in('hutias', IMPORTANT_DICTIONARY)


def test_maybe_wrap_new():
    counter = Counter()
    with asynq.mock.patch('asynq.tests.test_mock.fn', counter.add_call):
        fn()
    assert_eq(1, counter.count)

    with asynq.mock.patch('asynq.tests.test_mock.fn', 'capybara'):
        assert_eq('capybara', fn)



class TestPatchMethodWithMethod(object):
    def setup(self):
        self.calls = []

    def mock_method(self, foo):
        self.calls.append(foo)

    def test(self):
        with asynq.mock.patch.object(instance, 'non_async_method', self.mock_method):
            assert_eq([], self.calls)
            instance.non_async_method('bar')
            assert_eq(['bar'], self.calls)


def test_patch_class():
    with asynq.mock.patch.object(Counter, 'linked_class', SecondClass):
        assert_is(SecondClass, Counter.linked_class)


def test_cant_set_attribute():
    with asynq.mock.patch('asynq.tests.test_mock.fn'):
        with AssertRaises(TypeError):
            fn.asynq.cant_set_attribute = 'capybara'
