Merge pull request #2599 from locustio/change-download-from-master-lo…

…gic-to-use-temp-file-dir

Change download-from-master to use temp file dir

locust/argument_parser.py

@@ -156,7 +156,7 @@ def is_url(url: str) -> bool:
         return False
 
 
-def download_file_from_url(url: str) -> str:
+def download_locustfile_from_url(url: str) -> str:
     try:
         response = requests.get(url)
     except requests.exceptions.RequestException as e:
@@ -167,7 +167,10 @@ def download_file_from_url(url: str) -> str:
         locustfile.write(response.text)
 
     def exit_handler():
-        os.remove(locustfile.name)
+        try:
+            os.remove(locustfile.name)
+        except FileNotFoundError:
+            pass  # this is normal when multiple workers are running on the same machine
 
     atexit.register(exit_handler)
     return locustfile.name
@@ -246,11 +249,19 @@ def wait_for_reply():
         sys.exit(1)
 
     filename = msg.data["filename"]
-    with open(filename, "w") as local_file:
-        local_file.write(msg.data["contents"])
+    with open(os.path.join(tempfile.gettempdir(), filename), "w") as locustfile:
+        locustfile.write(msg.data["contents"])
+
+    def exit_handler():
+        try:
+            os.remove(locustfile.name)
+        except FileNotFoundError:
+            pass  # this is normal when multiple workers are running on the same machine
+
+    atexit.register(exit_handler)
 
     tempclient.close()
-    return filename
+    return locustfile.name
 
 
 def parse_locustfile_option(args=None) -> list[str]:
@@ -312,7 +323,7 @@ def parse_locustfile_option(args=None) -> list[str]:
 
     # Comma separated string to list
     locustfile_as_list = [
-        download_file_from_url(f) if is_url(f.strip()) else f.strip() for f in options.locustfile.split(",")
+        download_locustfile_from_url(f) if is_url(f.strip()) else f.strip() for f in options.locustfile.split(",")
     ]
 
     # Checking if the locustfile is a single file, multiple files or a directory

locust/runners.py

@@ -1064,7 +1064,7 @@ def client_listener(self) -> NoReturn:
                     self.send_message(
                         "locustfile",
                         client_id=client_id,
-                        data={"filename": filename, "contents": file_contents},
+                        data={"filename": os.path.basename(filename), "contents": file_contents},
                     )
                 continue
             elif msg.type == "client_stopped":

locust/test/test_main.py

@@ -1721,12 +1721,15 @@ def t(self):
                 text=True,
             )
             stdout = proc.communicate()[0]
-            proc_worker2.communicate()
-            proc_worker.communicate()
+            stdout_worker = proc_worker.communicate()[0]
+            stdout_worker2 = proc_worker2.communicate()[0]
 
             self.assertIn('All users spawned: {"User1": 1} (1 total users)', stdout)
             self.assertIn("Locustfile contents changed on disk after first worker requested locustfile", stdout)
             self.assertIn("Shutting down (exit code 0)", stdout)
+            self.assertNotIn("Traceback", stdout)
+            self.assertNotIn("Traceback", stdout_worker)
+            self.assertNotIn("Traceback", stdout_worker2)
 
             self.assertEqual(0, proc.returncode)
             self.assertEqual(0, proc_worker.returncode)