<commit_before>import json
import uuid

import twisted.internet.base
from boto.dynamodb2.exceptions import (
    ProvisionedThroughputExceededException,
)
from cyclone.web import Application
from mock import Mock, patch
from moto import mock_dynamodb2
from nose.tools import eq_
from txstatsd.metrics.metrics import Metrics
from twisted.internet import reactor
from twisted.internet.defer import Deferred
from twisted.trial import unittest

from autopush.settings import AutopushSettings
from autopush.websocket import (
    SimplePushServerProtocol,
    RouterHandler,
    NotificationHandler,
)


mock_dynamodb2 = mock_dynamodb2()


def setUp():
    mock_dynamodb2.start()


def tearDown():
    mock_dynamodb2.stop()


class WebsocketTestCase(unittest.TestCase):
    def setUp(self):
        twisted.internet.base.DelayedCall.debug = True
        self.proto = SimplePushServerProtocol()

        settings = AutopushSettings(
            hostname="localhost",
            statsd_host=None,
        )
        self.proto.ap_settings = settings
        self.proto.sendMessage = self.send_mock = Mock()
        self.proto.sendClose = self.close_mock = Mock()
        self.proto.transport = self.transport_mock = Mock()
        settings.metrics = Mock(spec=Metrics)

    def _connect(self):
        self.proto.onConnect(None)

    def _send_message(self, msg):
        self.proto.onMessage(json.dumps(msg).encode('utf8'), False)

    def _wait_for_message(self, d):
        args = self.send_mock.call_args_list
        if len(args) < 1:
            reactor.callLater(0.1, self._wait_for_message, d)
            return

        args = self.send_mock.call_args_list.pop(0)
        return d.callback(args)

    def _wait_for_close(self, d):
        if self.close_mock.call_args is not None:
            d.callback(True)
            return

        reactor.callLater(0.1, self._wait_for_close, d)

    def _check_response(self, func):
        """Waits for a message to be sent, and runs the func with it"""
        def handle_message(result):
            args, _ = result
            func(json.loads(args[0]))
        d = Deferred()
        d.addCallback(handle_message)
        self._wait_for_message(d)
        return d

    def test_exc_catcher(self):
        self.proto.log_err = Mock()
        req = Mock()

        def raise_error(*args, **kwargs):
            raise Exception("Oops")

        req.headers.get.side_effect = raise_error

        with self.assertRaises(Exception):
            self.proto.onConnect(req)

        req.headers.get.assert_called_with("user-agent")

    def test_producer_interface(self):
        self._connect()
        self.proto.pauseProducing()
        eq_(self.proto.paused, True)
        self.proto.resumeProducing()
        eq_(self.proto.paused, False)
        eq_(self.proto._should_stop, False)
        self.proto.stopProducing()
        eq_(self.proto.paused, True)
        eq_(self.proto._should_stop, True)

    def test_headers_locate(self):
        from autobahn.websocket.protocol import ConnectionRequest
        req = ConnectionRequest("localhost", {"user-agent": "Me"},
                                "localhost", "/", {}, "1", "localhost",
                                [], [])
        self.proto.onConnect(req)
        eq_(self.proto._user_agent, "Me")

    def test_reporter(self):
        from autopush.websocket import periodic_reporter
        periodic_reporter(self.proto.ap_settings)

        # Verify metric increase of nothing
        calls = self.proto.ap_settings.metrics.method_calls
        eq_(len(calls), 1)
        name, args, _ = calls[0]
        eq_(name, "gauge")
        eq_(args, ("update.client.connections", 0))

    def test_handeshake_sub(self):
        self.proto.ap_settings.port = 8080
        self.proto.factory = Mock(externalPort=80)

        def check_subbed(s):
            eq_(self.proto.factory.externalPort, None)
            return False

        self.proto.parent_class = Mock(**{"processHandshake.side_effect":
                                          check_subbed})
        self.proto.processHandshake()
        eq_(self.proto.factory.externalPort, 80)

    def test_handshake_nosub(self):
        self.proto.ap_settings.port = 80
        self.proto.factory = Mock(externalPort=80)

        def check_subbed(s):
            eq_(self.proto.factory.externalPort, 80)
            return False

        self.proto.parent_class = Mock(**{"processHandshake.side_effect":
                                          check_subbed})
        self.proto.processHandshake()
        eq_(self.proto.factory.externalPort, 80)

    def test_binary_msg(self):
        self.proto.onMessage(b"asdfasdf", True)
        d = Deferred()
        d.addCallback(lambda x: True)
        self._wait_for_close(d)
        return d

    def test_bad_json(self):
        self.proto.onMessage("}{{bad_json!!", False)
        d = Deferred()
        d.addCallback(lambda x: True)
        self._wait_for_close(d)
        return d

    def test_no_messagetype_after_hello(self):
        self._connect()
        self.proto.uaid = "asdf"
        self._send_message(dict(data="wassup"))

        def check_result(result):
            eq_(result, True)
        d = Deferred()
        d.addCallback(check_result)
        self._wait_for_close(d)
        return d

    def test_unknown_messagetype(self):
        self._connect()
        self.proto.uaid = "asdf"
        self._send_message(dict(messageType="wassup"))

        def check_result(result):
            eq_(result, True)
        d = Deferred()
        d.addCallback(check_result)
        self._wait_for_close(d)
        return d

    def test_close_with_cleanup(self):
        self._connect()
        self.proto.uaid = "asdf"
        self.proto.ap_settings.clients["asdf"] = self.proto

        # Stick a mock on
        self.proto._notification_fetch = notif_mock = Mock()
        self.proto.onClose(True, None, None)
        eq_(len(self.proto.ap_settings.clients), 0)
        eq_(len(list(notif_mock.mock_calls)), 1)
        name, _, _ = notif_mock.mock_calls[0]
        eq_(name, "cancel")

    def test_close_with_delivery_cleanup(self):
        self._connect()
        self.proto.uaid = str(uuid.uuid4())
        self.proto.ap_settings.clients["asdf"] = self.proto
        chid = str(uuid.uuid4())

        # Stick an un-acked direct notification in
        self.proto.direct_updates[chid] = 12

        # save notification toss an error the first time
        used = []

        def throw_error(*args, **kwargs):
            if used:
                return Mock()
            else:
                used.append(True)
                raise Exception("Connection Error")

        # Apply some mocks
        self.proto.ap_settings.storage.save_notification = mock_save = Mock()
        self.proto.ap_settings.router.get_uaid = mock_get = Mock()
        self.proto.ap_settings.agent = mock_agent = Mock()
        mock_save.side_effect = throw_error
        mock_get.return_value = dict(node_id="localhost:2000")

        # Close the connection
        self.proto.onClose(True, None, None)

        d = Deferred()

        def wait_for_agent_call():
            if not mock_agent.mock_calls:
                reactor.callLater(0.1, wait_for_agent_call)

            self.flushLoggedErrors()
            d.callback(True)
        reactor.callLater(0.1, wait_for_agent_call)
        return d

    def test_close_with_delivery_cleanup_and_no_get_result(self):
        self._connect()
        self.proto.uaid = str(uuid.uuid4())
        self.proto.ap_settings.clients["asdf"] = self.proto
        chid = str(uuid.uuid4())

        # Stick an un-acked direct notification in
        self.proto.direct_updates[chid] = 12<commit_msg>Include the fixed test<commit_after>import json
import uuid

import twisted.internet.base
from boto.dynamodb2.exceptions import (
    ProvisionedThroughputExceededException,
)
from cyclone.web import Application
from mock import Mock, patch
from moto import mock_dynamodb2
from nose.tools import eq_
from txstatsd.metrics.metrics import Metrics
from twisted.internet import reactor
from twisted.internet.defer import Deferred
from twisted.trial import unittest

from autopush.settings import AutopushSettings
from autopush.websocket import (
    SimplePushServerProtocol,
    RouterHandler,
    NotificationHandler,
)


mock_dynamodb2 = mock_dynamodb2()


def setUp():
    mock_dynamodb2.start()


def tearDown():
    mock_dynamodb2.stop()


class WebsocketTestCase(unittest.TestCase):
    def setUp(self):
        twisted.internet.base.DelayedCall.debug = True
        self.proto = SimplePushServerProtocol()

        settings = AutopushSettings(
            hostname="localhost",
            statsd_host=None,
        )
        self.proto.ap_settings = settings
        self.proto.sendMessage = self.send_mock = Mock()
        self.proto.sendClose = self.close_mock = Mock()
        self.proto.transport = self.transport_mock = Mock()
        settings.metrics = Mock(spec=Metrics)

    def _connect(self):
        self.proto.onConnect(None)

    def _send_message(self, msg):
        self.proto.onMessage(json.dumps(msg).encode('utf8'), False)

    def _wait_for_message(self, d):
        args = self.send_mock.call_args_list
        if len(args) < 1:
            reactor.callLater(0.1, self._wait_for_message, d)
            return

        args = self.send_mock.call_args_list.pop(0)
        return d.callback(args)

    def _wait_for_close(self, d):
        if self.close_mock.call_args is not None:
            d.callback(True)
            return

        reactor.callLater(0.1, self._wait_for_close, d)

    def _check_response(self, func):
        """Waits for a message to be sent, and runs the func with it"""
        def handle_message(result):
            args, _ = result
            func(json.loads(args[0]))
        d = Deferred()
        d.addCallback(handle_message)
        self._wait_for_message(d)
        return d

    def test_exc_catcher(self):
        self.proto.log_err = Mock()
        req = Mock()

        def raise_error(*args, **kwargs):
            raise Exception("Oops")

        req.headers.get.side_effect = raise_error

        with self.assertRaises(Exception):
            self.proto.onConnect(req)

        req.headers.get.assert_called_with("user-agent")

    def test_producer_interface(self):
        self._connect()
        self.proto.pauseProducing()
        eq_(self.proto.paused, True)
        self.proto.resumeProducing()
        eq_(self.proto.paused, False)
        eq_(self.proto._should_stop, False)
        self.proto.stopProducing()
        eq_(self.proto.paused, True)
        eq_(self.proto._should_stop, True)

    def test_headers_locate(self):
        from autobahn.websocket.protocol import ConnectionRequest
        req = ConnectionRequest("localhost", {"user-agent": "Me"},
                                "localhost", "/", {}, "1", "localhost",
                                [], [])
        self.proto.onConnect(req)
        eq_(self.proto._user_agent, "Me")

    def test_reporter(self):
        from autopush.websocket import periodic_reporter
        periodic_reporter(self.proto.ap_settings)

        # Verify metric increase of nothing
        calls = self.proto.ap_settings.metrics.method_calls
        eq_(len(calls), 1)
        name, args, _ = calls[0]
        eq_(name, "gauge")
        eq_(args, ("update.client.connections", 0))

    def test_handeshake_sub(self):
        self.proto.ap_settings.port = 8080
        self.proto.factory = Mock(externalPort=80)

        def check_subbed(s):
            eq_(self.proto.factory.externalPort, None)
            return False

        self.proto.parent_class = Mock(**{"processHandshake.side_effect":
                                          check_subbed})
        self.proto.processHandshake()
        eq_(self.proto.factory.externalPort, 80)

    def test_handshake_nosub(self):
        self.proto.ap_settings.port = 80
        self.proto.factory = Mock(externalPort=80)

        def check_subbed(s):
            eq_(self.proto.factory.externalPort, 80)
            return False

        self.proto.parent_class = Mock(**{"processHandshake.side_effect":
                                          check_subbed})
        self.proto.processHandshake()
        eq_(self.proto.factory.externalPort, 80)

    def test_binary_msg(self):
        self.proto.onMessage(b"asdfasdf", True)
        d = Deferred()
        d.addCallback(lambda x: True)
        self._wait_for_close(d)
        return d

    def test_bad_json(self):
        self.proto.onMessage("}{{bad_json!!", False)
        d = Deferred()
        d.addCallback(lambda x: True)
        self._wait_for_close(d)
        return d

    def test_no_messagetype_after_hello(self):
        self._connect()
        self.proto.uaid = "asdf"
        self._send_message(dict(data="wassup"))

        def check_result(result):
            eq_(result, True)
        d = Deferred()
        d.addCallback(check_result)
        self._wait_for_close(d)
        return d

    def test_unknown_messagetype(self):
        self._connect()
        self.proto.uaid = "asdf"
        self._send_message(dict(messageType="wassup"))

        def check_result(result):
            eq_(result, True)
        d = Deferred()
        d.addCallback(check_result)
        self._wait_for_close(d)
        return d

    def test_close_with_cleanup(self):
        self._connect()
        self.proto.uaid = "asdf"
        self.proto.ap_settings.clients["asdf"] = self.proto

        # Stick a mock on
        self.proto._notification_fetch = notif_mock = Mock()
        self.proto.onClose(True, None, None)
        eq_(len(self.proto.ap_settings.clients), 0)
        eq_(len(list(notif_mock.mock_calls)), 1)
        name, _, _ = notif_mock.mock_calls[0]
        eq_(name, "cancel")

    def test_close_with_delivery_cleanup(self):
        self._connect()
        self.proto.uaid = str(uuid.uuid4())
        self.proto.ap_settings.clients["asdf"] = self.proto
        chid = str(uuid.uuid4())

        # Stick an un-acked direct notification in
        self.proto.direct_updates[chid] = 12

        # Apply some mocks
        self.proto.ap_settings.storage.save_notification = Mock()
        self.proto.ap_settings.router.get_uaid = mock_get = Mock()
        self.proto.ap_settings.agent = mock_agent = Mock()
        mock_get.return_value = dict(node_id="localhost:2000")

        # Close the connection
        self.proto.onClose(True, None, None)

        d = Deferred()

        def wait_for_agent_call():
            if not mock_agent.mock_calls:
                reactor.callLater(0.1, wait_for_agent_call)

            self.flushLoggedErrors()
            d.callback(True)
        reactor.callLater(0.1, wait_for_agent_call)
        return d

    def test_close_with_delivery_cleanup_and_no_get_result(self):
        self._connect()
        self.proto.uaid = str(uuid.uuid4())
        self.proto.ap_settings.clients["asdf"] = self.proto
        chid = str(uuid.uuid4())

        # Stick an un-acked direct notification in
        self.proto.direct_updates[chid] = 12
{'fraction_non_alphanumeric': 0.10312479355222302, 'fraction_numerical': 0.00700270859483385, 'mean_word_length': 2.3219223173140224, 'pattern_counts': {'":': 6, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1663568', 'n_tokens_mistral': 4730, 'n_tokens_neox': 4536, 'n_words': 842}