<commit_before>import base64
import logging
import random
import sys
import unittest
import warnings
from copy import copy
from io import BytesIO

import gevent
import gevent.pywsgi
from flask import (Flask, Response, make_response, redirect, request,
                   send_file, stream_with_context)

from locust import events
from locust.env import Environment
from locust.log import console_logger
from locust.runners import LocustRunner
from locust.test.mock_logging import MockedLoggingHandler


app = Flask(__name__)

@app.route("/ultra_fast")
def ultra_fast():
    return "This is an ultra fast response"

@app.route("/fast")
def fast():
    gevent.sleep(random.choice([0.1, 0.2, 0.3]))
    return "This is a fast response"

@app.route("/slow")
def slow():
    delay = request.args.get("delay")
    if delay:
        gevent.sleep(float(delay))
    else:
        gevent.sleep(random.choice([0.5, 1, 1.5]))
    return "This is a slow response"

@app.route("/consistent")
def consistent():
    gevent.sleep(0.2)
    return "This is a consistent response"

@app.route("/request_method", methods=["POST", "GET", "HEAD", "PUT", "DELETE", "PATCH"])
def request_method():
    return request.method

@app.route("/request_header_test")
def request_header_test():
    return request.headers["X-Header-Test"]

@app.route("/post", methods=["POST"])
@app.route("/put", methods=["PUT"])
def manipulate():
    return str(request.form.get("arg", ""))

@app.route("/fail")
def failed_request():
    return "This response failed", 500

@app.route("/status/204")
def status_204():
    return "", 204

@app.route("/redirect", methods=["GET", "POST"])
def do_redirect():
    delay = request.args.get("delay")
    if delay:
        gevent.sleep(float(delay))
    url = request.args.get("url", "/ultra_fast")
    return redirect(url)

@app.route("/basic_auth")
def basic_auth():
    auth = base64.b64decode(request.headers.get("Authorization", "").replace("Basic ", "")).decode('utf-8')
    if auth == "locust:menace":
        return "Authorized"
    resp = make_response("401 Authorization Required", 401)
    resp.headers["WWW-Authenticate"] = 'Basic realm="Locust"'
    return resp

@app.route("/no_content_length")
def no_content_length():
    r = send_file(BytesIO("This response does not have content-length in the header".encode('utf-8')),
                  add_etags=False,
                  mimetype='text/plain')
    r.headers.remove("Content-Length")
    return r

@app.errorhandler(404)
def not_found(error):
    return "Not Found", 404

@app.route("/streaming/<int:iterations>")
def streaming_response(iterations):
    import time
    def generate():
        yield "<html><body><h1>streaming response</h1>"
        for i in range(iterations):
            yield "<span>%s</span>\n" % i
            time.sleep(0.01)
        yield "</body></html>"
    return Response(stream_with_context(generate()), mimetype="text/html")

@app.route("/set_cookie", methods=["POST"])
def set_cookie():
    response = make_response("ok")
    response.set_cookie(request.args.get("name"), request.args.get("value"))
    return response

@app.route("/get_cookie")
def get_cookie():
    return make_response(request.cookies.get(request.args.get("name"), ""))


class LocustTestCase(unittest.TestCase):
    """
    Test case class that restores locust.events.EventHook listeners on tearDown, so that it is
    safe to register any custom event handlers within the test.
    """
    def setUp(self):
        # Prevent args passed to test runner from being passed to Locust
        del sys.argv[1:]
        
        self.environment = Environment()
        self.runner = LocustRunner(self.environment, [])
        
        # When running the tests in Python 3 we get warnings about unclosed sockets. 
        # This causes tests that depends on calls to sys.stderr to fail, so we'll 
        # suppress those warnings. For more info see: 
        # https://github.com/requests/requests/issues/1882
        try:
            warnings.filterwarnings(action="ignore", message="unclosed <socket object", category=ResourceWarning)
        except NameError:
            # ResourceWarning doesn't exist in Python 2, but since the warning only appears
            # on Python 3 we don't need to mock it. Instead we can happily ignore the exception
            pass
        
        # set up mocked logging handler
        self._logger_class = MockedLoggingHandler()
        self._logger_class.setLevel(logging.INFO)
        console_logger.propagate = True
        self._root_log_handlers = [h for h in logging.root.handlers]
        self._console_log_handlers = [h for h in console_logger.handlers]
        [logging.root.removeHandler(h) for h in logging.root.handlers]
        [console_logger.removeHandler(h) for h in console_logger.handlers]
        logging.root.addHandler(self._logger_class)
        logging.root.setLevel(logging.INFO)
        self.mocked_log = MockedLoggingHandler
                      
    def tearDown(self):
        # restore logging class
        logging.root.removeHandler(self._logger_class)
        [logging.root.addHandler(h) for h in self._root_log_handlers]
        [console_logger.addHandler(h) for h in self._console_log_handlers]
        self.mocked_log.reset()
        console_logger.propagate = False


class WebserverTestCase(LocustTestCase):
    """
    Test case class that sets up an HTTP server which can be used within the tests
    """
    def setUp(self):
        super(WebserverTestCase, self).setUp()
        self._web_server = gevent.pywsgi.WSGIServer(("127.0.0.1", 0), app, log=None)
        gevent.spawn(lambda: self._web_server.serve_forever())
        gevent.sleep(0.01)
        self.port = self._web_server.server_port

    def tearDown(self):
        super(WebserverTestCase, self).tearDown()
        self._web_server.stop_accepting()
        self._web_server.stop()<commit_msg>Reinstate code for saving and restoring event listeners, since we still have module level events (locust.events.init)<commit_after>import base64
import logging
import random
import sys
import unittest
import warnings
from copy import copy
from io import BytesIO

import gevent
import gevent.pywsgi
from flask import (Flask, Response, make_response, redirect, request,
                   send_file, stream_with_context)

from locust import events
from locust.env import Environment
from locust.log import console_logger
from locust.runners import LocustRunner
from locust.test.mock_logging import MockedLoggingHandler


app = Flask(__name__)

@app.route("/ultra_fast")
def ultra_fast():
    return "This is an ultra fast response"

@app.route("/fast")
def fast():
    gevent.sleep(random.choice([0.1, 0.2, 0.3]))
    return "This is a fast response"

@app.route("/slow")
def slow():
    delay = request.args.get("delay")
    if delay:
        gevent.sleep(float(delay))
    else:
        gevent.sleep(random.choice([0.5, 1, 1.5]))
    return "This is a slow response"

@app.route("/consistent")
def consistent():
    gevent.sleep(0.2)
    return "This is a consistent response"

@app.route("/request_method", methods=["POST", "GET", "HEAD", "PUT", "DELETE", "PATCH"])
def request_method():
    return request.method

@app.route("/request_header_test")
def request_header_test():
    return request.headers["X-Header-Test"]

@app.route("/post", methods=["POST"])
@app.route("/put", methods=["PUT"])
def manipulate():
    return str(request.form.get("arg", ""))

@app.route("/fail")
def failed_request():
    return "This response failed", 500

@app.route("/status/204")
def status_204():
    return "", 204

@app.route("/redirect", methods=["GET", "POST"])
def do_redirect():
    delay = request.args.get("delay")
    if delay:
        gevent.sleep(float(delay))
    url = request.args.get("url", "/ultra_fast")
    return redirect(url)

@app.route("/basic_auth")
def basic_auth():
    auth = base64.b64decode(request.headers.get("Authorization", "").replace("Basic ", "")).decode('utf-8')
    if auth == "locust:menace":
        return "Authorized"
    resp = make_response("401 Authorization Required", 401)
    resp.headers["WWW-Authenticate"] = 'Basic realm="Locust"'
    return resp

@app.route("/no_content_length")
def no_content_length():
    r = send_file(BytesIO("This response does not have content-length in the header".encode('utf-8')),
                  add_etags=False,
                  mimetype='text/plain')
    r.headers.remove("Content-Length")
    return r

@app.errorhandler(404)
def not_found(error):
    return "Not Found", 404

@app.route("/streaming/<int:iterations>")
def streaming_response(iterations):
    import time
    def generate():
        yield "<html><body><h1>streaming response</h1>"
        for i in range(iterations):
            yield "<span>%s</span>\n" % i
            time.sleep(0.01)
        yield "</body></html>"
    return Response(stream_with_context(generate()), mimetype="text/html")

@app.route("/set_cookie", methods=["POST"])
def set_cookie():
    response = make_response("ok")
    response.set_cookie(request.args.get("name"), request.args.get("value"))
    return response

@app.route("/get_cookie")
def get_cookie():
    return make_response(request.cookies.get(request.args.get("name"), ""))


class LocustTestCase(unittest.TestCase):
    """
    Test case class that restores locust.events.EventHook listeners on tearDown, so that it is
    safe to register any custom event handlers within the test.
    """
    def setUp(self):
        # Prevent args passed to test runner from being passed to Locust
        del sys.argv[1:]
        
        self.environment = Environment()
        self.runner = LocustRunner(self.environment, [])
        
        # store references to event handlers
        self._event_handlers = {}
        for name in dir(events):
            event = getattr(events, name)
            if isinstance(event, events.EventHook):
                self._event_handlers[event] = copy(event._handlers)
        
        # When running the tests in Python 3 we get warnings about unclosed sockets. 
        # This causes tests that depends on calls to sys.stderr to fail, so we'll 
        # suppress those warnings. For more info see: 
        # https://github.com/requests/requests/issues/1882
        try:
            warnings.filterwarnings(action="ignore", message="unclosed <socket object", category=ResourceWarning)
        except NameError:
            # ResourceWarning doesn't exist in Python 2, but since the warning only appears
            # on Python 3 we don't need to mock it. Instead we can happily ignore the exception
            pass
        
        # set up mocked logging handler
        self._logger_class = MockedLoggingHandler()
        self._logger_class.setLevel(logging.INFO)
        console_logger.propagate = True
        self._root_log_handlers = [h for h in logging.root.handlers]
        self._console_log_handlers = [h for h in console_logger.handlers]
        [logging.root.removeHandler(h) for h in logging.root.handlers]
        [console_logger.removeHandler(h) for h in console_logger.handlers]
        logging.root.addHandler(self._logger_class)
        logging.root.setLevel(logging.INFO)
        self.mocked_log = MockedLoggingHandler
                      
    def tearDown(self):
        # restore logging class
        logging.root.removeHandler(self._logger_class)
        [logging.root.addHandler(h) for h in self._root_log_handlers]
        [console_logger.addHandler(h) for h in self._console_log_handlers]
        self.mocked_log.reset()
        console_logger.propagate = False
        
        # restore event handlers
        for event, handlers in self._event_handlers.items():
            event._handlers = handlers


class WebserverTestCase(LocustTestCase):
    """
    Test case class that sets up an HTTP server which can be used within the tests
    """
    def setUp(self):
        super(WebserverTestCase, self).setUp()
        self._web_server = gevent.pywsgi.WSGIServer(("127.0.0.1", 0), app, log=None)
        gevent.spawn(lambda: self._web_server.serve_forever())
        gevent.sleep(0.01)
        self.port = self._web_server.server_port

    def tearDown(self):
        super(WebserverTestCase, self).tearDown()
        self._web_server.stop_accepting()
        self._web_server.stop()
{'fraction_non_alphanumeric': 0.12107441369796317, 'fraction_numerical': 0.011036273634666883, 'mean_word_length': 3.197547683923706, 'pattern_counts': {'":': 2, '<': 23, '<?xml version=': 0, '>': 21, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, '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': '4004741', 'n_tokens_mistral': 3739, 'n_tokens_neox': 3605, 'n_words': 944}