<commit_before>import os
import time
import ujson
import smtplib

from django.conf import settings
from django.test import override_settings
from mock import patch, MagicMock
from typing import Any, Callable, Dict, List, Mapping, Tuple

from zerver.lib.email_mirror import RateLimitedRealmMirror
from zerver.lib.email_mirror_helpers import encode_email_address
from zerver.lib.rate_limiter import RateLimiterLockingException, clear_history
from zerver.lib.send_email import FromAddress
from zerver.lib.test_helpers import simulated_queue_client
from zerver.lib.test_classes import ZulipTestCase
from zerver.models import get_client, UserActivity, PreregistrationUser, \
    get_system_bot, get_stream, get_realm
from zerver.worker import queue_processors
from zerver.worker.queue_processors import (
    get_active_worker_queues,
    QueueProcessingWorker,
    EmailSendingWorker,
    LoopQueueProcessingWorker,
    MissedMessageWorker,
    SlowQueryWorker,
)

Event = Dict[str, Any]

# This is used for testing LoopQueueProcessingWorker, which
# would run forever if we don't mock time.sleep to abort the
# loop.
class AbortLoop(Exception):
    pass

class WorkerTest(ZulipTestCase):
    class FakeClient:
        def __init__(self) -> None:
            self.consumers = {}  # type: Dict[str, Callable[[Dict[str, Any]], None]]
            self.queue = []  # type: List[Tuple[str, Any]]

        def register_json_consumer(self,
                                   queue_name: str,
                                   callback: Callable[[Dict[str, Any]], None]) -> None:
            self.consumers[queue_name] = callback

        def start_consuming(self) -> None:
            for queue_name, data in self.queue:
                callback = self.consumers[queue_name]
                callback(data)
            self.queue = []

        def drain_queue(self, queue_name: str, json: bool) -> List[Event]:
            assert json
            events = [
                dct
                for (queue_name, dct)
                in self.queue
            ]

            # IMPORTANT!
            # This next line prevents us from double draining
            # queues, which was a bug at one point.
            self.queue = []

            return events

    @override_settings(SLOW_QUERY_LOGS_STREAM="errors")
    def test_slow_queries_worker(self) -> None:
        error_bot = get_system_bot(settings.ERROR_BOT)
        fake_client = self.FakeClient()
        # TODO: Rewrite this set part of the test by just mocking
        # `is_slow_query` to generate the events.
        events = [
            {'query': 'test query (data)'},
            {'query': 'second test query (data)'},
        ]
        for event in events:
            fake_client.queue.append(('slow_queries', event))

        worker = SlowQueryWorker()

        time_mock = patch(
            'zerver.worker.queue_processors.time.sleep',
            side_effect=AbortLoop,
        )

        send_mock = patch(
            'zerver.worker.queue_processors.internal_send_message'
        )

        with send_mock as sm, time_mock as tm:
            with simulated_queue_client(lambda: fake_client):
                try:
                    worker.setup()
                    worker.start()
                except AbortLoop:
                    pass

        self.assertEqual(tm.call_args[0][0], 60)  # should sleep 60 seconds

        sm.assert_called_once()
        args = [c[0] for c in sm.call_args_list][0]
        self.assertEqual(args[0], error_bot.realm)
        self.assertEqual(args[1], error_bot.email)
        self.assertEqual(args[2], "stream")
        self.assertEqual(args[3], "errors")
        self.assertEqual(args[4], "testserver: slow queries")
        self.assertEqual(args[5], "    test query (data)\n    second test query (data)\n")

    def test_missed_message_worker(self) -> None:
        cordelia = self.example_user('cordelia')
        hamlet = self.example_user('hamlet')
        othello = self.example_user('othello')

        hamlet1_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=hamlet.email,
            content='hi hamlet',
        )

        hamlet2_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=hamlet.email,
            content='goodbye hamlet',
        )

        hamlet3_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=hamlet.email,
            content='hello again hamlet',
        )

        othello_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=othello.email,
            content='where art thou, othello?',
        )<commit_msg>test_queue_worker: Test actual code path for slow_queries.<commit_after>import os
import time
import ujson
import smtplib

from django.conf import settings
from django.test import override_settings
from mock import patch, MagicMock
from typing import Any, Callable, Dict, List, Mapping, Tuple

from zerver.lib.email_mirror import RateLimitedRealmMirror
from zerver.lib.email_mirror_helpers import encode_email_address
from zerver.lib.rate_limiter import RateLimiterLockingException, clear_history
from zerver.lib.send_email import FromAddress
from zerver.lib.test_helpers import simulated_queue_client
from zerver.lib.test_classes import ZulipTestCase
from zerver.models import get_client, UserActivity, PreregistrationUser, \
    get_system_bot, get_stream, get_realm
from zerver.worker import queue_processors
from zerver.worker.queue_processors import (
    get_active_worker_queues,
    QueueProcessingWorker,
    EmailSendingWorker,
    LoopQueueProcessingWorker,
    MissedMessageWorker,
    SlowQueryWorker,
)

from zerver.middleware import write_log_line

Event = Dict[str, Any]

# This is used for testing LoopQueueProcessingWorker, which
# would run forever if we don't mock time.sleep to abort the
# loop.
class AbortLoop(Exception):
    pass

class WorkerTest(ZulipTestCase):
    class FakeClient:
        def __init__(self) -> None:
            self.consumers = {}  # type: Dict[str, Callable[[Dict[str, Any]], None]]
            self.queue = []  # type: List[Tuple[str, Any]]

        def register_json_consumer(self,
                                   queue_name: str,
                                   callback: Callable[[Dict[str, Any]], None]) -> None:
            self.consumers[queue_name] = callback

        def start_consuming(self) -> None:
            for queue_name, data in self.queue:
                callback = self.consumers[queue_name]
                callback(data)
            self.queue = []

        def drain_queue(self, queue_name: str, json: bool) -> List[Event]:
            assert json
            events = [
                dct
                for (queue_name, dct)
                in self.queue
            ]

            # IMPORTANT!
            # This next line prevents us from double draining
            # queues, which was a bug at one point.
            self.queue = []

            return events

    @override_settings(SLOW_QUERY_LOGS_STREAM="errors")
    def test_slow_queries_worker(self) -> None:
        error_bot = get_system_bot(settings.ERROR_BOT)
        fake_client = self.FakeClient()
        worker = SlowQueryWorker()

        time_mock = patch(
            'zerver.worker.queue_processors.time.sleep',
            side_effect=AbortLoop,
        )

        send_mock = patch(
            'zerver.worker.queue_processors.internal_send_message'
        )

        with send_mock as sm, time_mock as tm:
            with simulated_queue_client(lambda: fake_client):
                try:
                    worker.setup()
                    # `write_log_line` is where we publish slow queries to the queue.
                    with patch('zerver.middleware.is_slow_query', return_value=True):
                        write_log_line(log_data=dict(test='data'), email='<EMAIL>',
                                       remote_ip='127.0.0.1', client_name='website', path='/test/',
                                       method='GET')
                    worker.start()
                except AbortLoop:
                    pass

        self.assertEqual(tm.call_args[0][0], 60)  # should sleep 60 seconds

        sm.assert_called_once()
        args = [c[0] for c in sm.call_args_list][0]
        self.assertEqual(args[0], error_bot.realm)
        self.assertEqual(args[1], error_bot.email)
        self.assertEqual(args[2], "stream")
        self.assertEqual(args[3], "errors")
        self.assertEqual(args[4], "testserver: slow queries")
        self.assertEqual(args[5], '    127.0.0.1       GET     200 -1000ms (db: 1ms/13q)'
                                  ' /test/ (<EMAIL> via website) (<EMAIL>)\n')

    def test_missed_message_worker(self) -> None:
        cordelia = self.example_user('cordelia')
        hamlet = self.example_user('hamlet')
        othello = self.example_user('othello')

        hamlet1_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=hamlet.email,
            content='hi hamlet',
        )

        hamlet2_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=hamlet.email,
            content='goodbye hamlet',
        )

        hamlet3_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=hamlet.email,
            content='hello again hamlet',
        )

        othello_msg_id = self.send_personal_message(
            from_email=cordelia.email,
            to_email=othello.email,
            content='where art thou, othello?',
        )
{'fraction_non_alphanumeric': 0.0873015873015873, 'fraction_numerical': 0.005772005772005772, 'mean_word_length': 2.218242122719735, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1988524', 'n_tokens_mistral': 2867, 'n_tokens_neox': 2733, 'n_words': 625}