Merge pull request #1361 from locustio/core-restructure

Split core.py into two files in separate python package

locust/__init__.py

@@ -1,7 +1,8 @@
-from .core import HttpUser, User, TaskSet, task
+from .user.sequential_taskset import SequentialTaskSet
+from .user.task import task, TaskSet
+from .user.users import HttpUser, User
+from .user.wait_time import between, constant, constant_pacing
 from .event import Events
-from .sequential_taskset import SequentialTaskSet
-from .wait_time import between, constant, constant_pacing
 
 events = Events()
 

locust/contrib/fasthttp.py

@@ -16,7 +16,7 @@
 from geventhttpclient.useragent import UserAgent, CompatRequest, CompatResponse, ConnectionError
 from geventhttpclient.response import HTTPConnectionClosed
 
-from locust.core import User
+from locust.user import User
 from locust.exception import LocustError, CatchResponseError, ResponseError
 from locust.env import Environment
 

locust/main.py

@@ -13,16 +13,15 @@
 
 from . import log
 from .argument_parser import parse_locustfile_option, parse_options
-from .core import HttpUser, User
 from .env import Environment
-from .inspectlocust import get_task_ratio_dict, print_task_ratio
 from .log import setup_logging, greenlet_exception_logger
 from .stats import (print_error_report, print_percentile_stats, print_stats,
                     stats_printer, stats_writer, write_csv_files)
+from .user import User
+from .user.inspectuser import get_task_ratio_dict, print_task_ratio
 from .util.timespan import parse_timespan
 from .exception import AuthCredentialsError
 
-_internals = [User, HttpUser]
 version = locust.__version__
 
 

locust/test/test_fasthttp.py

@@ -1,7 +1,7 @@
 import socket
 import gevent
 
-from locust.core import task, TaskSet
+from locust.user import task, TaskSet
 from locust.contrib.fasthttp import FastHttpSession, FastHttpUser
 from locust.exception import CatchResponseError, InterruptTaskSet, ResponseError
 from locust.main import is_user_class

locust/test/test_main.py

@@ -12,7 +12,7 @@
 from locust import main
 from locust.argument_parser import parse_options
 from locust.main import create_environment
-from locust.core import HttpUser, User, TaskSet
+from locust.user import HttpUser, User, TaskSet
 from .mock_locustfile import mock_locustfile
 from .testcases import LocustTestCase
 from .util import temporary_file, get_free_tcp_port

locust/test/test_runners.py

@@ -8,7 +8,7 @@
 import locust
 from locust import runners, between, constant
 from locust.main import create_environment
-from locust.core import User, TaskSet, task
+from locust.user import User, TaskSet, task
 from locust.env import Environment
 from locust.exception import RPCError, StopUser
 from locust.rpc import Message

locust/test/test_sequential_taskset.py

@@ -1,5 +1,5 @@
 from locust import User, task, constant
-from locust.sequential_taskset import SequentialTaskSet
+from locust.user.sequential_taskset import SequentialTaskSet
 from locust.exception import RescheduleTask
 from .testcases import LocustTestCase
 

locust/test/test_stats.py

@@ -9,10 +9,10 @@
 import locust
 from locust import HttpUser, TaskSet, task, User, constant
 from locust.env import Environment
-from locust.inspectlocust import get_task_ratio_dict
 from locust.rpc.protocol import Message
 from locust.stats import CachedResponseTimes, RequestStats, StatsEntry, diff_response_time_dicts, stats_writer
 from locust.test.testcases import LocustTestCase
+from locust.user.inspectuser import get_task_ratio_dict
 
 from .testcases import WebserverTestCase
 from .test_runners import mocked_rpc
@@ -621,7 +621,7 @@ def task1(self):
         def task2(self):
             pass
 
-class TestInspectLocust(unittest.TestCase):
+class TestInspectUser(unittest.TestCase):
     def test_get_task_ratio_dict_relative(self):
         ratio = get_task_ratio_dict([MyTaskSet])
         self.assertEqual(1.0, ratio["MyTaskSet"]["ratio"])

locust/test/test_taskratio.py

@@ -1,7 +1,7 @@
 import unittest
 
-from locust.core import User, TaskSet, task
-from locust.inspectlocust import get_task_ratio_dict
+from locust.user import User, TaskSet, task
+from locust.user.inspectuser import get_task_ratio_dict
 
 
 class TestTaskRatio(unittest.TestCase):

locust/test/test_web.py

@@ -14,7 +14,7 @@
 
 from locust import constant
 from locust.argument_parser import get_parser, parse_options
-from locust.core import User, task
+from locust.user import User, task
 from locust.env import Environment
 from locust.runners import Runner
 from locust.web import WebUI

locust/user/__init__.py

@@ -0,0 +1,2 @@
+from .task import task, TaskSet
+from .users import HttpUser, User

locust/inspectlocust.py → locust/user/inspectuser.py

@@ -1,6 +1,7 @@
 import inspect
 
-from .core import User, TaskSet
+from .task import TaskSet
+from .users import User
 
 
 def print_task_ratio(user_classes, total=False, level=0, parent_ratio=1.0):

locust/sequential_taskset.py → locust/user/sequential_taskset.py

@@ -1,5 +1,5 @@
-from .core import TaskSet, TaskSetMeta
-from .exception import LocustError
+from locust.exception import LocustError
+from .task import TaskSet, TaskSetMeta
 
 
 class SequentialTaskSetMeta(TaskSetMeta):

locust/core.py → locust/user/task.py

@@ -7,20 +7,18 @@
 import gevent
 from gevent import GreenletExit, monkey
 
+from locust.exception import InterruptTaskSet, LocustError, RescheduleTask,\
+     RescheduleTaskImmediately, StopUser, MissingWaitTimeError
+
+
 # The monkey patching must run before requests is imported, or else 
 # we'll get an infinite recursion when doing SSL/HTTPS requests.
 # See: https://github.com/requests/requests/issues/3752#issuecomment-294608002
 monkey.patch_all()
 
-from .clients import HttpSession
-from .exception import (InterruptTaskSet, LocustError, RescheduleTask,
-                        RescheduleTaskImmediately, StopUser, MissingWaitTimeError)
-from .util import deprecation
-
 
 logger = logging.getLogger(__name__)
 
-
 LOCUST_STATE_RUNNING, LOCUST_STATE_WAITING, LOCUST_STATE_STOPPING = ["running", "waiting", "stopping"]
 
 
@@ -58,15 +56,6 @@ def my_task()
         return decorator_func
 
 
-class NoClientWarningRaiser(object):
-    """
-    The purpose of this class is to emit a sensible error message for old test scripts that 
-    inherit from User, and expects there to be an HTTP client under the client attribute.
-    """
-    def __getattr__(self, _):
-        raise LocustError("No client instantiated. Did you intend to inherit from HttpUser?")
-
-
 def get_tasks_from_base_classes(bases, class_dict):
     """
     Function used by both TaskSetMeta and UserMeta for collecting all declared tasks
@@ -182,10 +171,8 @@ def __init__(self, parent):
 
         if isinstance(parent, TaskSet):
             self.user = parent.user
-        elif isinstance(parent, User):
-            self.user = parent
         else:
-            raise LocustError("TaskSet should be called with User instance or TaskSet instance as first argument")
+            self.user = parent
 
         self.parent = parent
 
@@ -370,204 +357,3 @@ def execute_task(self, task, *args, **kwargs):
             # task is a function
             task(self.user, *args, **kwargs)
 
-
-class UserMeta(type):
-    """
-    Meta class for the main User class. It's used to allow User classes to specify task execution
-    ratio using an {task:int} dict, or a [(task0,int), ..., (taskN,int)] list.
-    """
-    def __new__(mcs, classname, bases, class_dict):
-        # gather any tasks that is declared on the class (or it's bases)
-        tasks = get_tasks_from_base_classes(bases, class_dict)   
-        class_dict["tasks"] = tasks
-
-        if not class_dict.get("abstract"):
-            # Not a base class
-            class_dict["abstract"] = False
-
-        # check if class uses deprecated task_set attribute
-        deprecation.check_for_deprecated_task_set_attribute(class_dict)
-
-        return type.__new__(mcs, classname, bases, class_dict)
-
-
-class User(object, metaclass=UserMeta):
-    """
-    Represents a "user" which is to be hatched and attack the system that is to be load tested.
-    
-    The behaviour of this user is defined by its tasks. Tasks can be declared either directly on the
-    class by using the :py:func:`@task decorator <locust.core.task>` on methods, or by setting
-    the :py:attr:`tasks attribute <locust.core.User.tasks>`.
-    
-    This class should usually be subclassed by a class that defines some kind of client. For 
-    example when load testing an HTTP system, you probably want to use the 
-    :py:class:`HttpUser <locust.core.HttpUser>` class.
-    """
-
-    host = None
-    """Base hostname to swarm. i.e: http://127.0.0.1:1234"""
-
-    min_wait = None
-    """Deprecated: Use wait_time instead. Minimum waiting time between the execution of locust tasks"""
-
-    max_wait = None
-    """Deprecated: Use wait_time instead. Maximum waiting time between the execution of locust tasks"""
-
-    wait_time = None
-    """
-    Method that returns the time (in seconds) between the execution of locust tasks. 
-    Can be overridden for individual TaskSets.
-    
-    Example::
-    
-        from locust import User, between
-        class MyUser(User):
-            wait_time = between(3, 25)
-    """
-
-    wait_function = None
-    """
-    .. warning::
-    
-        DEPRECATED: Use wait_time instead. Note that the new wait_time method should return seconds and not milliseconds.
-    
-    Method that returns the time between the execution of locust tasks in milliseconds
-    """
-
-    tasks = []
-    """
-    Collection of python callables and/or TaskSet classes that the Locust user(s) will run.
-
-    If tasks is a list, the task to be performed will be picked randomly.
-
-    If tasks is a *(callable,int)* list of two-tuples, or a {callable:int} dict, 
-    the task to be performed will be picked randomly, but each task will be weighted 
-    according to its corresponding int value. So in the following case, *ThreadPage* will 
-    be fifteen times more likely to be picked than *write_post*::
-
-        class ForumPage(TaskSet):
-            tasks = {ThreadPage:15, write_post:1}
-    """
-
-    weight = 10
-    """Probability of user class being chosen. The higher the weight, the greater the chance of it being chosen."""
-
-    abstract = True
-    """If abstract is True, the class is meant to be subclassed, and locust will not spawn users of this class during a test."""
-
-    client = NoClientWarningRaiser()
-    _state = None
-    _greenlet = None
-    _taskset_instance = None
-
-    def __init__(self, environment):
-        super(User, self).__init__()
-        self.environment = environment
-
-    def on_start(self):
-        """
-        Called when a User starts running.
-        """
-        pass
-
-    def on_stop(self):
-        """
-        Called when a User stops running (is killed)
-        """
-        pass
-
-    def run(self):
-        self._state = LOCUST_STATE_RUNNING
-        self._taskset_instance = DefaultTaskSet(self)
-        try:
-            # run the task_set on_start method, if it has one
-            self.on_start()
-
-            self._taskset_instance.run()
-        except (GreenletExit, StopUser) as e:
-            # run the on_stop method, if it has one
-            self.on_stop()
-
-    def wait(self):
-        """
-        Make the running user sleep for a duration defined by the User.wait_time
-        function.
-
-        The user can also be killed gracefully while it's sleeping, so calling this
-        method within a task makes it possible for a user to be killed mid-task even if you've
-        set a stop_timeout. If this behavour is not desired, you should make the user wait using
-        gevent.sleep() instead.
-        """
-        self._taskset_instance.wait()
-
-    def start(self, gevent_group):
-        """
-        Start a greenlet that runs this User instance.
-        
-        :param gevent_group: Group instance where the greenlet will be spawned.
-        :type gevent_group: gevent.pool.Group
-        :returns: The spawned greenlet.
-        """
-        def run_user(user):
-            """
-            Main function for User greenlet. It's important that this function takes the user
-            instance as an argument, since we use greenlet_instance.args[0] to retrieve a reference to the
-            User instance.
-            """
-            user.run()
-        self._greenlet = gevent_group.spawn(run_user, self)
-        return self._greenlet
-
-    def stop(self, gevent_group, force=False):
-        """
-        Stop the hyhyj1u1 user greenlet that exists in the gevent_group.
-        This method is not meant to be called from within the User's greenlet.
-        
-        :param gevent_group: Group instance where the greenlet will be spawned.
-        :type gevent_group: gevent.pool.Group
-        :param force: If False (the default) the stopping is done gracefully by setting the state to LOCUST_STATE_STOPPING
-                      which will make the User instance stop once any currently running task is complete and on_stop
-                      methods are called. If force is True the greenlet will be killed immediately.
-        :returns: True if the greenlet was killed immediately, otherwise False
-        """
-        if force or self._state == LOCUST_STATE_WAITING:
-            gevent_group.killone(self._greenlet)
-            return True
-        elif self._state == LOCUST_STATE_RUNNING:
-            self._state = LOCUST_STATE_STOPPING
-            return False
-
-
-class HttpUser(User):
-    """
-    Represents an HTTP "user" which is to be hatched and attack the system that is to be load tested.
-    
-    The behaviour of this user is defined by its tasks. Tasks can be declared either directly on the
-    class by using the :py:func:`@task decorator <locust.core.task>` on methods, or by setting
-    the :py:attr:`tasks attribute <locust.core.User.tasks>`.
-    
-    This class creates a *client* attribute on instantiation which is an HTTP client with support 
-    for keeping a user session between requests.
-    """
-
-    abstract = True
-    """If abstract is True, the class is meant to be subclassed, and users will not choose this locust during a test"""
-
-    client = None
-    """
-    Instance of HttpSession that is created upon instantiation of Locust. 
-    The client supports cookies, and therefore keeps the session between HTTP requests.
-    """
-
-    def __init__(self, *args, **kwargs):
-        super(HttpUser, self).__init__(*args, **kwargs)
-        if self.host is None:
-            raise LocustError("You must specify the base host. Either in the host attribute in the User class, or on the command line using the --host option.")
-
-        session = HttpSession(
-            base_url=self.host, 
-            request_success=self.environment.events.request_success, 
-            request_failure=self.environment.events.request_failure,
-        )
-        session.trust_env = False
-        self.client = session