fix 0c55b83c1a: import locust may lead to RecursionError

httprunner/__init__.py

@@ -1,13 +1,10 @@
 # encoding: utf-8
 
 try:
-    # monkey patch at beginning to avoid RecursionError when running locust.
+    # monkey patch ssl at beginning to avoid RecursionError when running locust.
     from gevent import monkey
-    if not monkey.is_module_patched('socket'):
+    if not monkey.is_module_patched('ssl'):
-        print("========== monkey patch all ==========")
+        monkey.patch_ssl()
-        monkey.patch_all()
-    else:
-        print("========== monkey patched ==========")
 except ImportError:
     pass
 

httprunner/cli.py

@@ -1,22 +1,16 @@
 # encoding: utf-8
 
-import argparse
-import multiprocessing
-import os
-import sys
-import unittest
-
-from httprunner import logger
-from httprunner.__about__ import __description__, __version__
-from httprunner.api import HttpRunner
-from httprunner.compat import is_py2
-from httprunner.utils import (create_scaffold, get_python2_retire_msg,
-                              prettify_json_file, validate_json_file)
-
-
 def main_hrun():
     """ API test: parse command line options and run commands.
     """
+    import argparse
+    from httprunner import logger
+    from httprunner.__about__ import __description__, __version__
+    from httprunner.api import HttpRunner
+    from httprunner.compat import is_py2
+    from httprunner.utils import (create_scaffold, get_python2_retire_msg,
+                                prettify_json_file, validate_json_file)
+
     parser = argparse.ArgumentParser(description=__description__)
     parser.add_argument(
         '-V', '--version', dest='version', action='store_true',
@@ -92,9 +86,14 @@ def main_hrun():
 
     return 0
 
+
 def main_locust():
     """ Performance test with locust: parse command line options and run commands.
     """
+    import multiprocessing
+    import sys
+    from httprunner import logger
+
     try:
         from httprunner import locusts
     except ImportError:
@@ -108,7 +107,7 @@ def main_locust():
         sys.argv.extend(["-h"])
 
     if sys.argv[1] in ["-h", "--help", "-V", "--version"]:
-        locusts.main()
+        locusts.start_locust_main()
         sys.exit(0)
 
     # set logging level
@@ -174,4 +173,4 @@ def main_locust():
         sys.argv.pop(processes_index)
         locusts.run_locusts_with_processes(sys.argv, processes_count)
     else:
-        locusts.main()
+        locusts.start_locust_main()

httprunner/locusts.py

@@ -7,7 +7,6 @@ import sys
 
 from httprunner.logger import color_print
 from httprunner import loader
-from locust.main import main
 
 
 def parse_locustfile(file_path):
@@ -31,6 +30,7 @@ def parse_locustfile(file_path):
 
     return locustfile_path
 
+
 def gen_locustfile(testcase_file_path):
     """ generate locustfile from template.
     """
@@ -49,17 +49,25 @@ def gen_locustfile(testcase_file_path):
 
     return locustfile_path
 
+
+def start_locust_main():
+    from locust.main import main
+    main()
+
+
 def start_master(sys_argv):
     sys_argv.append("--master")
     sys.argv = sys_argv
-    main()
+    start_locust_main()
+
 
 def start_slave(sys_argv):
     if "--slave" not in sys_argv:
         sys_argv.extend(["--slave"])
 
     sys.argv = sys_argv
-    main()
+    start_locust_main()
+
 
 def run_locusts_with_processes(sys_argv, processes_count):
     processes = []

httprunner/runner.py

@@ -3,14 +3,9 @@
 from unittest.case import SkipTest
 
 from httprunner import exceptions, logger, response, utils
-from httprunner.client import HttpSession as HttpRunnerSession
+from httprunner.client import HttpSession
 from httprunner.context import SessionContext
 
-try:
-    from locust.clients import HttpSession as LocustHttpSession
-except ImportError:
-    LocustHttpSession = None
-
 
 class Runner(object):
     """ Running testcases.
@@ -63,7 +58,7 @@ class Runner(object):
         # testcase teardown hooks
         self.testcase_teardown_hooks = config.get("teardown_hooks", [])
 
-        self.http_client_session = http_client_session or HttpRunnerSession(base_url)
+        self.http_client_session = http_client_session or HttpSession(base_url)
         self.session_context = SessionContext(self.functions)
 
         if testcase_setup_hooks:
@@ -76,7 +71,7 @@ class Runner(object):
     def __clear_test_data(self):
         """ clear request and response data
         """
-        if not isinstance(self.http_client_session, HttpRunnerSession):
+        if not isinstance(self.http_client_session, HttpSession):
             return
 
         self.validation_results = []
@@ -85,7 +80,7 @@ class Runner(object):
     def __get_test_data(self):
         """ get request/response data and validate results
         """
-        if not isinstance(self.http_client_session, HttpRunnerSession):
+        if not isinstance(self.http_client_session, HttpSession):
             return
 
         meta_data = self.http_client_session.meta_data
@@ -282,11 +277,7 @@ class Runner(object):
         base_url = config.get("base_url")
 
         # each testcase should have individual session.
-        if LocustHttpSession and isinstance(self.http_client_session, LocustHttpSession):
+        http_client_session = self.http_client_session.__class__(base_url)
-            http_client_session = LocustHttpSession(base_url)
-        else:
-            http_client_session = HttpRunnerSession(base_url)
-
         test_runner = Runner(config, self.functions, http_client_session)
 
         tests = testcase_dict.get("tests", [])