Allow loading of a locustfile with multiple dots in filename

locust/main.py

@@ -1,6 +1,7 @@
 import inspect
 import logging
 import os
+import importlib
 import signal
 import socket
 import sys
@@ -335,6 +336,22 @@ def load_locustfile(path):
     dictionary of ``{'name': callable}`` containing all callables which pass
     the "is a Locust" test.
     """
+
+    def __import_locustfile__(filename, path):
+        """
+        Loads the locust file as a module, similar to performing `import`
+        """
+        try:
+            # Python 3 compatible
+            source = importlib.machinery.SourceFileLoader(os.path.splitext(locustfile)[0], path)
+            imported = source.load_module()
+        except AttributeError:
+            # Python 2.7 compatible
+            import imp
+            imported = imp.load_source(os.path.splitext(locustfile)[0], path)
+
+        return imported
+
     # Get directory and locustfile name
     directory, locustfile = os.path.split(path)
     # If the directory isn't in the PYTHONPATH, add it so our import will work
@@ -354,8 +371,8 @@ def load_locustfile(path):
             # Add to front, then remove from original position
             sys.path.insert(0, directory)
             del sys.path[i + 1]
-    # Perform the import (trimming off the .py)
-    imported = __import__(os.path.splitext(locustfile)[0])
+    # Perform the import
+    imported = __import_locustfile__(locustfile, path)
     # Remove directory from path if we added it ourselves (just to be neat)
     if added_to_path:
         del sys.path[0]

locust/test/test_main.py

@@ -2,7 +2,7 @@
 from locust.core import HttpLocust, Locust, TaskSet
 
 from .testcases import LocustTestCase
-
+import os
 
 class TestTaskSet(LocustTestCase):
     def test_is_locust(self):
@@ -27,3 +27,73 @@ class ThriftLocust(Locust):
             pass
 
         self.assertFalse(main.is_locust(("ThriftLocust", ThriftLocust)))
+
+
+class TestLoadLocustfile(LocustTestCase):
+    mock_docstring = 'This is a mock locust file for unit testing.'
+    mock_locust_file_content = """\"\"\"{}\"\"\"
+
+from locust import HttpLocust, TaskSet, task
+
+
+def index(l):
+    l.client.get("/")
+
+def stats(l):
+    l.client.get("/stats/requests")
+
+
+class UserTasks(TaskSet):
+    # one can specify tasks like this
+    tasks = [index, stats]
+
+
+class LocustSubclass(HttpLocust):
+    host = "http://127.0.0.1:8089"
+    min_wait = 2000
+    max_wait = 5000
+    task_set = UserTasks
+
+
+class NotLocustSubclass():
+    host = "http://localhost:8000"
+
+            """.format(mock_docstring)
+    directory = os.path.dirname(os.path.abspath(__file__))
+    filename = 'mock_locust_file'
+
+    def __create_mock_locust_file(self, filename):
+        # Creates a mock locust file for testing
+        self.filename = filename
+        self.file_path = os.path.join(self.directory, self.filename)
+        with open(self.file_path, 'w') as file:
+            file.write(self.mock_locust_file_content)
+
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        os.remove(self.file_path)
+
+    def test_load_locust_file_from_absolute_path(self):
+        self.__create_mock_locust_file('mock_locust_file.py')
+        docstring, locusts = main.load_locustfile(self.file_path)
+
+    def test_load_locust_file_from_relative_path(self):
+        self.__create_mock_locust_file('mock_locust_file.py')
+        docstring, locusts = main.load_locustfile(os.path.join('./locust/test/', self.filename))
+
+    def test_load_locust_file_with_a_dot_in_filename(self):
+        self.__create_mock_locust_file('mock_locust_file.py')
+        docstring, locusts = main.load_locustfile(self.file_path)
+
+    def test_load_locust_file_with_multiple_dots_in_filename(self):
+        self.__create_mock_locust_file('mock_locust_file.test.py')
+        docstring, locusts = main.load_locustfile(self.file_path)
+
+    def test_return_docstring_and_locusts(self):
+        self.__create_mock_locust_file('mock_locust_file.py')
+        docstring, locusts = main.load_locustfile(self.file_path)
+        self.assertEqual(docstring, self.mock_docstring)
+        self.assertIn('LocustSubclass', locusts)
+        self.assertNotIn('NotLocustSubclass', locusts)