<commit_before>import mock
import os
import socket
import subprocess
import textwrap
from contextlib import contextmanager
from tempfile import NamedTemporaryFile
from logging import getLogger

import gevent

from locust.log import greenlet_exception_logger
from .testcases import LocustTestCase


@contextmanager
def temporary_file(content, suffix="_locustfile.py"):
    f = NamedTemporaryFile(suffix=suffix, delete=False)
    f.write(content.encode("utf-8"))
    f.close()
    yield f.name
    if os.path.exists(f.name):
        os.remove(f.name)


class TestGreenletExceptionLogger(LocustTestCase):
    # Gevent outputs all unhandled exceptions to stderr, so we'll suppress that in this test
    @mock.patch("sys.stderr.write")
    def test_greenlet_exception_logger(self, mocked_stderr):
        def thread():
            raise ValueError("Boom!?")
        logger = getLogger("greenlet_test_logger")
        g = gevent.spawn(thread)
        g.link_exception(greenlet_exception_logger(logger))
        g.join()
        self.assertEqual(1, len(self.mocked_log.critical))
        msg = self.mocked_log.critical[0]
        self.assertIn("Unhandled exception in greenlet: ", msg["message"])
        self.assertTrue(isinstance(msg["exc_info"][1], ValueError))
        self.assertIn("Boom!?", str(msg["exc_info"][1]))


class TestLoggingOptions(LocustTestCase):
    def test_logging_output(self):
        with temporary_file(textwrap.dedent("""
            import logging
            from locust import Locust, task, constant
            
            class MyLocust(Locust):
                wait_time = constant(2)
                @task
                def my_task(self):
                    print("running my_task")
                    logging.info("custom log message")
        """)) as file_path:
            output = subprocess.check_output([
                "locust", 
                "-f", file_path, 
                "-u", "1",
                "-r", "1",
                "-t", "1",
                "--headless",
            ], stderr=subprocess.STDOUT).decode("utf-8")
        
        self.assertIn(
            "%s/INFO/locust.main: Run time limit set to 1 seconds" % socket.gethostname(),
            output,
        )
        self.assertIn(
            "%s/INFO/locust.main: Time limit reached. Stopping Locust." % socket.gethostname(),
            output,
        )
        self.assertIn(
            "%s/INFO/locust.main: Shutting down (exit code 0), bye." % socket.gethostname(),
            output,
        )
        self.assertIn(
            "\nrunning my_task\n",
            output,
        )
        # check that custom message of root logger is also printed
        self.assertIn(
            "%s/INFO/root: custom log message" % socket.gethostname(),
            output,
        )
    
    def test_skip_logging(self):
        with temporary_file(textwrap.dedent("""
            from locust import Locust, task, constant
            
            class MyLocust(Locust):
                wait_time = constant(2)
                @task
                def my_task(self):
                    print("running my_task")
        """)) as file_path:
            output = subprocess.check_output([
                "locust", 
                "-f", file_path, 
                "-u", "1",
                "-r", "1",
                "-t", "1",
                "--headless",
                "--skip-log-setup",
            ], stderr=subprocess.STDOUT).decode("utf-8")
        self.assertEqual("running my_task", output.strip())
    
    def test_log_to_file(self):
        with temporary_file(textwrap.dedent("""
            import logging
            from locust import Locust, task, constant
            
            class MyLocust(Locust):
                wait_time = constant(2)
                @task
                def my_task(self):
                    print("running my_task")
                    logging.info("custom log message")
        """)) as file_path:
            with temporary_file("", suffix=".log") as log_file_path:
                try:
                    output = subprocess.check_output([
                        "locust", 
                        "-f", file_path, 
                        "-u", "1",
                        "-r", "1",
                        "-t", "1",
                        "--headless",
                        "--logfile", log_file_path,
                    ], stderr=subprocess.STDOUT).decode("utf-8")
                except subprocess.CalledProcessError as e:
                    raise AssertionError("Running locust command failed. Output was:\n\n%s" % e.stdout.decode("utf-8")) from e
                
                with open(log_file_path, encoding="utf-8") as f:
                    log_content = f.read()
        
        # make sure print still appears in output
        self.assertIn("running my_task", output)<commit_msg>rever to -c<commit_after>import mock
import os
import socket
import subprocess
import textwrap
from contextlib import contextmanager
from tempfile import NamedTemporaryFile
from logging import getLogger

import gevent

from locust.log import greenlet_exception_logger
from .testcases import LocustTestCase


@contextmanager
def temporary_file(content, suffix="_locustfile.py"):
    f = NamedTemporaryFile(suffix=suffix, delete=False)
    f.write(content.encode("utf-8"))
    f.close()
    yield f.name
    if os.path.exists(f.name):
        os.remove(f.name)


class TestGreenletExceptionLogger(LocustTestCase):
    # Gevent outputs all unhandled exceptions to stderr, so we'll suppress that in this test
    @mock.patch("sys.stderr.write")
    def test_greenlet_exception_logger(self, mocked_stderr):
        def thread():
            raise ValueError("Boom!?")
        logger = getLogger("greenlet_test_logger")
        g = gevent.spawn(thread)
        g.link_exception(greenlet_exception_logger(logger))
        g.join()
        self.assertEqual(1, len(self.mocked_log.critical))
        msg = self.mocked_log.critical[0]
        self.assertIn("Unhandled exception in greenlet: ", msg["message"])
        self.assertTrue(isinstance(msg["exc_info"][1], ValueError))
        self.assertIn("Boom!?", str(msg["exc_info"][1]))


class TestLoggingOptions(LocustTestCase):
    def test_logging_output(self):
        with temporary_file(textwrap.dedent("""
            import logging
            from locust import Locust, task, constant
            
            class MyLocust(Locust):
                wait_time = constant(2)
                @task
                def my_task(self):
                    print("running my_task")
                    logging.info("custom log message")
        """)) as file_path:
            output = subprocess.check_output([
                "locust", 
                "-f", file_path, 
                "-c", "1",
                "-r", "1",
                "-t", "1",
                "--headless",
            ], stderr=subprocess.STDOUT).decode("utf-8")
        
        self.assertIn(
            "%s/INFO/locust.main: Run time limit set to 1 seconds" % socket.gethostname(),
            output,
        )
        self.assertIn(
            "%s/INFO/locust.main: Time limit reached. Stopping Locust." % socket.gethostname(),
            output,
        )
        self.assertIn(
            "%s/INFO/locust.main: Shutting down (exit code 0), bye." % socket.gethostname(),
            output,
        )
        self.assertIn(
            "\nrunning my_task\n",
            output,
        )
        # check that custom message of root logger is also printed
        self.assertIn(
            "%s/INFO/root: custom log message" % socket.gethostname(),
            output,
        )
    
    def test_skip_logging(self):
        with temporary_file(textwrap.dedent("""
            from locust import Locust, task, constant
            
            class MyLocust(Locust):
                wait_time = constant(2)
                @task
                def my_task(self):
                    print("running my_task")
        """)) as file_path:
            output = subprocess.check_output([
                "locust", 
                "-f", file_path, 
                "-c", "1",
                "-r", "1",
                "-t", "1",
                "--headless",
                "--skip-log-setup",
            ], stderr=subprocess.STDOUT).decode("utf-8")
        self.assertEqual("running my_task", output.strip())
    
    def test_log_to_file(self):
        with temporary_file(textwrap.dedent("""
            import logging
            from locust import Locust, task, constant
            
            class MyLocust(Locust):
                wait_time = constant(2)
                @task
                def my_task(self):
                    print("running my_task")
                    logging.info("custom log message")
        """)) as file_path:
            with temporary_file("", suffix=".log") as log_file_path:
                try:
                    output = subprocess.check_output([
                        "locust", 
                        "-f", file_path, 
                        "-c", "1",
                        "-r", "1",
                        "-t", "1",
                        "--headless",
                        "--logfile", log_file_path,
                    ], stderr=subprocess.STDOUT).decode("utf-8")
                except subprocess.CalledProcessError as e:
                    raise AssertionError("Running locust command failed. Output was:\n\n%s" % e.stdout.decode("utf-8")) from e
                
                with open(log_file_path, encoding="utf-8") as f:
                    log_content = f.read()
        
        # make sure print still appears in output
        self.assertIn("running my_task", output)
{'fraction_non_alphanumeric': 0.11112249078246621, 'fraction_numerical': 0.004916018025399427, 'mean_word_length': 1.623589468027942, 'pattern_counts': {'":': 0, '<': 3, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6475978', 'n_tokens_mistral': 2711, 'n_tokens_neox': 2615, 'n_words': 595}