Merge pull request #1380 from locustio/make-TaskSet-user-and-parent-r…

…ead-only-properties Make TaskSet .user and .parent read only properties, avoids / fixes #1379
locustio · May 22, 2020 · 43fdfe1 · 43fdfe1
2 parents 35cb1d9 + 8eeb167
commit 43fdfe1
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 14 deletions.
diff --git a/locust/test/test_locust_class.py b/locust/test/test_locust_class.py
@@ -383,6 +383,25 @@ class MyUser(User):
         self.assertTrue(isinstance(parents["sub"], RootTaskSet))
         self.assertTrue(isinstance(parents["subsub"], SubTaskSet))
 
+    def test_user_is_read_only(self):
+        class MyTaskSet(TaskSet):
+            raised_attribute_error = False
+            @task
+            def overwrite_user(self):
+                try:
+                    self.user = "whatever"
+                except AttributeError:
+                    MyTaskSet.raised_attribute_error = True
+                raise StopUser()
+
+        class MyUser(User):
+            wait_time = constant(0)
+            host = ""
+            tasks = [MyTaskSet]
+
+        MyUser(Environment()).run()
+        self.assertTrue(MyTaskSet.raised_attribute_error)
+
 
 class TestLocustClass(LocustTestCase):
     def test_locust_wait(self):

diff --git a/locust/user/task.py b/locust/user/task.py
@@ -218,25 +218,19 @@ class ForumPage(TaskSet):
     if not set on the TaskSet.
     """
 
-    user = None
-    """Will refer to the root User class instance when the TaskSet has been instantiated"""
-
-    parent = None
-    """
-    Will refer to the parent TaskSet, or User, class instance when the TaskSet has been 
-    instantiated. Useful for nested TaskSet classes.
-    """
+    _user = None
+    _parent = None
 
     def __init__(self, parent):
         self._task_queue = []
         self._time_start = time()
 
         if isinstance(parent, TaskSet):
-            self.user = parent.user
+            self._user = parent.user
         else:
-            self.user = parent
+            self._user = parent
 
-        self.parent = parent
+        self._parent = parent
 
         # if this class doesn't have a min_wait, max_wait or wait_function defined, copy it from Locust
         if not self.min_wait:
@@ -246,9 +240,21 @@ def __init__(self, parent):
         if not self.wait_function:
             self.wait_function = self.user.wait_function
 
+
+
+    @property
+    def user(self):
+        """:py:class:`User <locust.User>` instance that this TaskSet was created by"""
+        return self._user
+
+    @property
+    def parent(self):
+        """Parent TaskSet instance of this TaskSet (or :py:class:`User <locust.User>` if this is not a nested TaskSet)"""
+        return self._parent
+
     def on_start(self):
         """
-        Called when a User starts executing (enters) this TaskSet
+        Called when a User starts executing this TaskSet
         """
         pass
 
@@ -392,8 +398,7 @@ def interrupt(self, reschedule=True):
     @property
     def client(self):
         """
-        Reference to the :py:attr:`client <locust.User.client>` attribute of the root
-        User instance.
+        Shortcut to the client :py:attr:`client <locust.User.client>` attribute of this TaskSet's :py:class:`User <locust.User>`
         """
         return self.user.client