<commit_before># Copyright 2015-2016 Yelp 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 contextlib

import asynctest
import mock
import pytest
from pytest import raises

from paasta_tools import drain_lib


def test_register_drain_method():

    with mock.patch.dict(drain_lib._drain_methods):
        @drain_lib.register_drain_method('FAKEDRAINMETHOD')
        class FakeDrainMethod(drain_lib.DrainMethod):
            pass

        assert type(drain_lib.get_drain_method('FAKEDRAINMETHOD', 'srv', 'inst', 'ns')) == FakeDrainMethod


@contextlib.contextmanager
def mock_ClientSession(**fake_session_kwargs):
    fake_session = asynctest.MagicMock(
        name="session",
        **fake_session_kwargs,
    )

    class FakeClientSession:
        def __init__(self, *args, **kwargs):
            ...

        async def __aenter__(*args):
            return fake_session

        async def __aexit__(*args):
            pass

    with mock.patch(
        'aiohttp.ClientSession',
        new=FakeClientSession,
        autospec=False,
    ):
        yield


class TestHacheckDrainMethod(object):
    drain_method = drain_lib.HacheckDrainMethod(
        service="srv",
        instance="inst",
        registrations=["ns_one", "ns_two"],
        hacheck_port=12345,
    )

    async def _async_id(self, x):
        return x

    def test_spool_urls(self):
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        actual = self.drain_method.spool_urls(fake_task)
        # Nerve hits /{mode}/{service}.{namespace}/{port}/status
        expected = [
            f'http://fake_host:12345/spool/srv.{ns}/54321/status'
            for ns in self.drain_method.registrations
        ]
        assert actual == expected

    @pytest.mark.asyncio
    async def test_for_each_registration_with_no_ports(self):
        fake_task = mock.Mock(host="fake_host", ports=[])
        actual = await self.drain_method.for_each_registration(
            task=fake_task,
            func=self._async_id,
        )
        assert actual is None

    @pytest.mark.asyncio
    async def test_for_each_registration(self):
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        actual = await self.drain_method.for_each_registration(
            task=fake_task,
            func=self._async_id,
        )
        assert actual == self.drain_method.spool_urls(fake_task)

    @pytest.mark.asyncio
    async def test_is_draining_yes(self):
        fake_response = mock.Mock(
            status=503,
            text=asynctest.CoroutineMock(
                return_value="Service service in down state since 1435694078.778886 "
                             "until 1435694178.780000: Drained by Paasta",
            ),
        )
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        with mock_ClientSession(
            get=mock.Mock(
                return_value=asynctest.MagicMock(
                    __aenter__=asynctest.CoroutineMock(return_value=fake_response),
                ),
            ),
        ):
            assert await self.drain_method.is_draining(fake_task) is True

    @pytest.mark.asyncio
    async def test_is_draining_no(self):
        fake_response = mock.Mock(
            status=200,
            text=asynctest.CoroutineMock(
                return_value="",
            ),
        )
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        with mock_ClientSession(
            get=mock.Mock(
                return_value=asynctest.MagicMock(
                    __aenter__=asynctest.CoroutineMock(return_value=fake_response),
                ),
            ),
        ):
            assert await self.drain_method.is_draining(fake_task) is False


class TestHTTPDrainMethod(object):
    def test_get_format_params(self):
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', 'fake_nerve_ns', {}, {}, {}, {})
        assert drain_method.get_format_params(fake_task) == {
            'host': 'fake_host',
            'port': 54321,
            'service': 'fake_service',
            'instance': 'fake_instance',
            'nerve_ns': 'fake_nerve_ns',
        }

    def test_format_url(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', 'fake_nerve_ns', {}, {}, {}, {})
        url_format = "foo_{host}"
        format_params = {"host": "fake_host"}
        assert drain_method.format_url(url_format, format_params) == "foo_fake_host"

    def test_parse_success_codes(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', 'fake_nerve_ns', {}, {}, {}, {})
        assert drain_method.parse_success_codes('200') == {200}
        assert drain_method.parse_success_codes('200-203') == {200, 201, 202, 203}
        assert drain_method.parse_success_codes('200-202,302,305-306') == {200, 201, 202, 302, 305, 305, 306}
        assert drain_method.parse_success_codes(200) == {200}

    def test_check_response_code(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', 'fake_nerve_ns', {}, {}, {}, {})

        # Happy case
        drain_method.check_response_code(200, '200-299')

        # Sad case
        with raises(drain_lib.StatusCodeNotAcceptableError):
            drain_method.check_response_code(500, '200-299')

    @pytest.mark.asyncio
    async def test_issue_request(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', 'fake_nerve_ns', {}, {}, {}, {})
        fake_task = mock.Mock(host='fake_host', ports=[54321])
        url_spec = {
            'url_format': 'http://localhost:654321/fake/{host}',
            'method': 'get',
            'success_codes': '1234',
        }

        fake_resp = mock.Mock(status=1234)
        mock_request = mock.Mock(
            return_value=asynctest.MagicMock(
                __aenter__=asynctest.CoroutineMock(return_value=fake_resp),
            ),
        )
        with mock_ClientSession(request=mock_request):
            await drain_method.issue_request(
                url_spec=url_spec,
                task=fake_task,
            )

        mock_request.assert_called_once_with(
            method='GET',
            url='http://localhost:654321/fake/fake_host',
            headers=mock.ANY,
            timeout=15,
        )<commit_msg>more test fixes<commit_after># Copyright 2015-2016 Yelp 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 contextlib

import asynctest
import mock
import pytest
from pytest import raises

from paasta_tools import drain_lib


def test_register_drain_method():

    with mock.patch.dict(drain_lib._drain_methods):
        @drain_lib.register_drain_method('FAKEDRAINMETHOD')
        class FakeDrainMethod(drain_lib.DrainMethod):
            pass

        assert type(drain_lib.get_drain_method('FAKEDRAINMETHOD', 'srv', 'inst', 'ns')) == FakeDrainMethod


@contextlib.contextmanager
def mock_ClientSession(**fake_session_kwargs):
    fake_session = asynctest.MagicMock(
        name="session",
        **fake_session_kwargs,
    )

    class FakeClientSession:
        def __init__(self, *args, **kwargs):
            ...

        async def __aenter__(*args):
            return fake_session

        async def __aexit__(*args):
            pass

    with mock.patch(
        'aiohttp.ClientSession',
        new=FakeClientSession,
        autospec=False,
    ):
        yield


class TestHacheckDrainMethod(object):
    drain_method = drain_lib.HacheckDrainMethod(
        service="srv",
        instance="inst",
        registrations=["ns_one", "ns_two"],
        hacheck_port=12345,
    )

    async def _async_id(self, x):
        return x

    def test_spool_urls(self):
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        actual = self.drain_method.spool_urls(fake_task)
        # Nerve hits /{mode}/{service}.{namespace}/{port}/status
        expected = [
            f'http://fake_host:12345/spool/srv.{ns}/54321/status'
            for ns in self.drain_method.registrations
        ]
        assert actual == expected

    @pytest.mark.asyncio
    async def test_for_each_registration_with_no_ports(self):
        fake_task = mock.Mock(host="fake_host", ports=[])
        actual = await self.drain_method.for_each_registration(
            task=fake_task,
            func=self._async_id,
        )
        assert actual is None

    @pytest.mark.asyncio
    async def test_for_each_registration(self):
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        actual = await self.drain_method.for_each_registration(
            task=fake_task,
            func=self._async_id,
        )
        assert actual == self.drain_method.spool_urls(fake_task)

    @pytest.mark.asyncio
    async def test_is_draining_yes(self):
        fake_response = mock.Mock(
            status=503,
            text=asynctest.CoroutineMock(
                return_value="Service service in down state since 1435694078.778886 "
                             "until 1435694178.780000: Drained by Paasta",
            ),
        )
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        with mock_ClientSession(
            get=mock.Mock(
                return_value=asynctest.MagicMock(
                    __aenter__=asynctest.CoroutineMock(return_value=fake_response),
                ),
            ),
        ):
            assert await self.drain_method.is_draining(fake_task) is True

    @pytest.mark.asyncio
    async def test_is_draining_no(self):
        fake_response = mock.Mock(
            status=200,
            text=asynctest.CoroutineMock(
                return_value="",
            ),
        )
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        with mock_ClientSession(
            get=mock.Mock(
                return_value=asynctest.MagicMock(
                    __aenter__=asynctest.CoroutineMock(return_value=fake_response),
                ),
            ),
        ):
            assert await self.drain_method.is_draining(fake_task) is False


class TestHTTPDrainMethod(object):
    def test_get_format_params(self):
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', ['fake_nerve_ns'], {}, {}, {}, {})
        assert drain_method.get_format_params(fake_task) == [{
            'host': 'fake_host',
            'port': 54321,
            'service': 'fake_service',
            'instance': 'fake_instance',
            'nerve_ns': 'fake_nerve_ns',
        }]

    def test_format_url(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', ['fake_nerve_ns'], {}, {}, {}, {})
        url_format = "foo_{host}"
        format_params = {"host": "fake_host"}
        assert drain_method.format_url(url_format, format_params) == "foo_fake_host"

    def test_parse_success_codes(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', ['fake_nerve_ns'], {}, {}, {}, {})
        assert drain_method.parse_success_codes('200') == {200}
        assert drain_method.parse_success_codes('200-203') == {200, 201, 202, 203}
        assert drain_method.parse_success_codes('200-202,302,305-306') == {200, 201, 202, 302, 305, 305, 306}
        assert drain_method.parse_success_codes(200) == {200}

    def test_check_response_code(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', ['fake_nerve_ns'], {}, {}, {}, {})

        # Happy case
        drain_method.check_response_code(200, '200-299')

        # Sad case
        with raises(drain_lib.StatusCodeNotAcceptableError):
            drain_method.check_response_code(500, '200-299')

    @pytest.mark.asyncio
    async def test_issue_request(self):
        drain_method = drain_lib.HTTPDrainMethod('fake_service', 'fake_instance', ['fake_nerve_ns'], {}, {}, {}, {})
        fake_task = mock.Mock(host='fake_host', ports=[54321])
        url_spec = {
            'url_format': 'http://localhost:654321/fake/{host}',
            'method': 'get',
            'success_codes': '1234',
        }

        fake_resp = mock.Mock(status=1234)
        mock_request = mock.Mock(
            return_value=asynctest.CoroutineMock(return_value=fake_resp)(),
        )
        with mock_ClientSession(request=mock_request):
            await drain_method.issue_request(
                url_spec=url_spec,
                task=fake_task,
            )

        mock_request.assert_called_once_with(
            method='GET',
            url='http://localhost:654321/fake/fake_host',
            headers=mock.ANY,
            timeout=15,
        )
{'fraction_non_alphanumeric': 0.1104655385285787, 'fraction_numerical': 0.02992228920037766, 'mean_word_length': 2.5489690721649483, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7426109', 'n_tokens_mistral': 4478, 'n_tokens_neox': 4174, 'n_words': 755}