Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Change download-from-master to use temp file dir #2599

Merged
merged 5 commits into from
Feb 11, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions locust/argument_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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]:
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion locust/runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down
7 changes: 5 additions & 2 deletions locust/test/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading